fix: Transformers that return multiple items are not properly handled

nodestream/pipeline/transformers/transformer.py

@@ -1,5 +1,5 @@
 from abc import abstractmethod
-from typing import Any, AsyncGenerator, Generator
+from typing import Any, AsyncGenerator
 
 from ..flush import Flush
 from ..step import Step
@@ -19,12 +19,12 @@ async def handle_async_record_stream(
             if record is Flush:
                 yield record
             else:
-                val_or_gen = await self.transform_record(record)
-                if isinstance(val_or_gen, Generator):
-                    for result in val_or_gen:
+                val_or_gen = self.transform_record(record)
+                if isinstance(val_or_gen, AsyncGenerator):
+                    async for result in val_or_gen:
                         yield result
                 else:
-                    yield val_or_gen
+                    yield await val_or_gen
 
     @abstractmethod
     async def transform_record(self, record: Any) -> Any:

tests/unit/pipeline/transformers/test_expand_json_field.py

@@ -14,5 +14,9 @@
 @pytest.mark.parametrize("input,output,path", [(SIMPLE_INPUT, SIMPLE_OUTPUT, "b")])
 async def test_expand_json_fields(input, output, path):
     subject = ExpandJsonField.from_file_data(path)
-    result = await subject.transform_record(input)
-    assert_that(result, equal_to(output))
+
+    async def upstream():
+        yield input
+
+    results = [r async for r in subject.handle_async_record_stream(upstream())]
+    assert_that(results, equal_to([output]))

tests/unit/pipeline/transformers/test_value_projection.py

@@ -10,5 +10,9 @@ async def test_value_projection_transform_record():
     subject = ValueProjection(
         projection=JmespathValueProvider.from_string_expression("items[*]")
     )
-    results = [r async for r in subject.transform_record(record={"items": [1, 2, 3]})]
+
+    async def record():
+        yield {"items": [1, 2, 3]}
+
+    results = [r async for r in subject.handle_async_record_stream(record())]
     assert_that(results, equal_to([1, 2, 3]))