From 07347a7f0d5783570ff03cd842540261d560ec60 Mon Sep 17 00:00:00 2001
From: marcosschroh <schrohm@gmail.com>
Date: Wed, 6 Dec 2023 15:45:12 +0100
Subject: [PATCH] fix(typing): StreamFunc type to match the udf signature.

---
 kstreams/streams.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/kstreams/streams.py b/kstreams/streams.py
index 5d28af6..7e18513 100644
--- 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],
     *,