fix(typing): StreamFunc type to match the udf signature. (#142)

kpn · Dec 6, 2023 · 4a58047 · 4a58047
1 parent ee1756c
commit 4a58047
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/kstreams/streams.py b/kstreams/streams.py
@@ -30,6 +30,9 @@
 
 logger = logging.getLogger(__name__)
 
+# Function required by the `stream` decorator
+StreamFunc = Callable[..., Awaitable[Any]]
+
 
 class Stream:
     """
@@ -86,7 +89,7 @@ def __init__(
         self,
         topics: Union[List[str], str],
         *,
-        func: Callable[..., Awaitable[Any]],
+        func: StreamFunc,
         backend: Optional[Kafka] = None,
         consumer_class: Type[ConsumerType] = Consumer,
         name: Optional[str] = None,
@@ -272,10 +275,6 @@ async def __anext__(self) -> ConsumerRecord:
             raise StopAsyncIteration  # noqa: F821
 
 
-# Function required by the `stream` decorator
-StreamFunc = Callable[[Stream], Awaitable[Any]]
-
-
 def stream(
     topics: Union[List[str], str],
     *,