Refactor

Author: gtopper

storey/flow.py

@@ -26,7 +26,19 @@
 from asyncio import Task
 from collections import defaultdict
 from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
-from typing import Any, Callable, Collection, Dict, Iterable, List, Optional, Set, Union
+from typing import (
+    Any,
+    AsyncGenerator,
+    Callable,
+    Collection,
+    Dict,
+    Generator,
+    Iterable,
+    List,
+    Optional,
+    Set,
+    Union,
+)
 
 import aiohttp
 
@@ -45,6 +57,11 @@
 from .utils import _split_path, get_in, stringify_key, update_in
 
 
+def _is_generator(obj) -> bool:
+    """Check if an object is a sync or async generator."""
+    return inspect.isgenerator(obj) or inspect.isasyncgen(obj)
+
+
 class Flow:
     _legal_first_step = False
 
@@ -534,11 +551,6 @@ class _StreamingStepMixin:
     want to support user-provided generator functions.
     """
 
-    @staticmethod
-    def _is_generator(obj) -> bool:
-        """Check if an object is a sync or async generator."""
-        return inspect.isgenerator(obj) or inspect.isasyncgen(obj)
-
     def _validate_not_already_streaming(self, event):
         """Ensure we're not streaming on top of an already streaming event.
 
@@ -552,35 +564,35 @@ def _validate_not_already_streaming(self, event):
                 f"Step '{self.name}' received a streaming event from '{streaming_step}'."
             )
 
-    async def _emit_streaming_chunks(self, original_event, generator):
+    async def _emit_streaming_chunks(self, event, generator: Generator | AsyncGenerator) -> None:
         """Emit streaming chunks from a generator, then send StreamCompletion.
 
         Args:
-            original_event: The original event that triggered this streaming response.
+            event: The event that will be used to create chunk events.
             generator: A sync or async generator yielding chunk bodies.
         """
-        self._validate_not_already_streaming(original_event)
+        self._validate_not_already_streaming(event)
 
         chunk_id = 0
         if inspect.isgenerator(generator):
             # Sync generator
             for chunk_body in generator:
-                chunk_event = self._user_fn_output_to_event(original_event, chunk_body)
+                chunk_event = self._user_fn_output_to_event(event, chunk_body)
                 chunk_event.streaming_step = self.name
                 chunk_event.chunk_id = chunk_id
                 await self._do_downstream(chunk_event)
                 chunk_id += 1
         else:
             # Async generator
             async for chunk_body in generator:
-                chunk_event = self._user_fn_output_to_event(original_event, chunk_body)
+                chunk_event = self._user_fn_output_to_event(event, chunk_body)
                 chunk_event.streaming_step = self.name
                 chunk_event.chunk_id = chunk_id
                 await self._do_downstream(chunk_event)
                 chunk_id += 1
 
         # Send completion signal
-        await self._do_downstream(StreamCompletion(self.name, original_event))
+        await self._do_downstream(StreamCompletion(self.name, event))
 
 
 class _UnaryFunctionFlow(Flow):
@@ -671,7 +683,7 @@ class Map(_UnaryFunctionFlow, _StreamingStepMixin):
 
     async def _do_internal(self, event, fn_result):
         # Check if the result is a generator (streaming response)
-        if self._is_generator(fn_result):
+        if _is_generator(fn_result):
             await self._emit_streaming_chunks(event, fn_result)
         else:
             mapped_event = self._user_fn_output_to_event(event, fn_result)
@@ -864,7 +876,7 @@ async def _do(self, event):
         fn_result = await self._call(element)
         if not self._filter:
             # Check if the result is a generator (streaming response)
-            if self._is_generator(fn_result):
+            if _is_generator(fn_result):
                 await self._emit_streaming_chunks(event, fn_result)
             else:
                 mapped_event = self._user_fn_output_to_event(event, fn_result)
@@ -1806,18 +1818,13 @@ async def run_async(self, body: Any, path: str, origin_name: Optional[str] = Non
         """
         return body
 
-    @staticmethod
-    def _is_generator(obj) -> bool:
-        """Check if an object is a sync or async generator."""
-        return inspect.isgenerator(obj) or inspect.isasyncgen(obj)
-
     def _run(self, body: Any, path: str, origin_name: Optional[str] = None) -> Any:
         timestamp = datetime.datetime.now(tz=datetime.timezone.utc)
         start = time.monotonic()
         try:
             result = self.run(body, path, origin_name)
             # Return generator directly for streaming support
-            if self._is_generator(result):
+            if _is_generator(result):
                 return result
             body = result
         except Exception as e:
@@ -1837,7 +1844,7 @@ async def _async_run(self, body: Any, path: str, origin_name: Optional[str] = No
                 return self.run_async(body, path, origin_name)
             result = await self.run_async(body, path, origin_name)
             # Return generator directly for streaming support
-            if self._is_generator(result):
+            if _is_generator(result):
                 return result
             body = result
         except Exception as e:
@@ -2154,7 +2161,7 @@ async def _do(self, event):
         if len(runnables) == 1 and results:
             result = results[0]
             # Check if the result is a generator (streaming response)
-            if self._is_generator(result):
+            if _is_generator(result):
                 # Validate execution mechanism - streaming not supported with process-based execution
                 runnable_name = (
                     runnables[0].name if isinstance(runnables[0], ParallelExecutionRunnable) else runnables[0]
@@ -2181,7 +2188,7 @@ async def _do(self, event):
         else:
             # Check if any results are generators (not allowed with multiple runnables)
             for result in results:
-                if self._is_generator(result):
+                if _is_generator(result):
                     raise StreamingError(
                         "Streaming is not supported when multiple runnables are selected. "
                         "Streaming runnables must be the only runnable selected for an event."

tests/test_streaming.py

@@ -33,7 +33,7 @@
     build_flow,
 )
 from storey.dtypes import Event, StreamChunk, StreamCompletion
-from storey.flow import _StreamingStepMixin
+from storey.flow import _is_generator
 
 
 class TestStreamingPrimitives:
@@ -60,27 +60,27 @@ def test_stream_completion_repr(self):
         assert "my_step" in repr(completion)
 
 
-class TestStreamingStepMixin:
-    """Tests for the _StreamingStepMixin utility methods."""
+class TestIsGenerator:
+    """Tests for the _is_generator utility function."""
 
     def test_is_generator_sync(self):
         def gen():
             yield 1
             yield 2
 
-        assert _StreamingStepMixin._is_generator(gen())
+        assert _is_generator(gen())
 
     def test_is_generator_async(self):
         async def async_gen():
             yield 1
             yield 2
 
-        assert _StreamingStepMixin._is_generator(async_gen())
+        assert _is_generator(async_gen())
 
     def test_is_generator_non_generator(self):
-        assert not _StreamingStepMixin._is_generator([1, 2, 3])
-        assert not _StreamingStepMixin._is_generator("string")
-        assert not _StreamingStepMixin._is_generator(42)
+        assert not _is_generator([1, 2, 3])
+        assert not _is_generator("string")
+        assert not _is_generator(42)
 
     def test_is_generator_coroutine(self):
         async def coro():
@@ -89,7 +89,7 @@ async def coro():
         # Coroutine is not a generator
         c = coro()
         try:
-            assert not _StreamingStepMixin._is_generator(c)
+            assert not _is_generator(c)
         finally:
             c.close()