fix(typing): pyi, examples, and manifest (#309)

Signed-off-by: Vigith Maurice <vigith@gmail.com>

Author: vigith

packages/pynumaflow-lite/manifests/batchmap/batchmap_cat.py

@@ -1,6 +1,7 @@
 import asyncio
 import signal
 from collections.abc import AsyncIterable
+from typing import Awaitable, Callable
 
 from pynumaflow_lite import batchmapper
 from pynumaflow_lite.batchmapper import Message
@@ -28,7 +29,7 @@ async def handler(self, batch: AsyncIterable[batchmapper.Datum]) -> batchmapper.
     pass
 
 
-async def start(f: callable):
+async def start(f: Callable[[AsyncIterable[batchmapper.Datum]], Awaitable[batchmapper.BatchResponses]]):
     server = batchmapper.BatchMapAsyncServer()
 
     # Register loop-level signal handlers so we control shutdown and avoid asyncio.run

packages/pynumaflow-lite/manifests/map/map_cat.py

@@ -1,5 +1,7 @@
 import asyncio
 import signal
+from typing import Awaitable, Callable
+
 from pynumaflow_lite import mapper
 
 
@@ -26,7 +28,7 @@ async def handler(
     pass
 
 
-async def start(f: callable):
+async def start(f: Callable[[list[str], mapper.Datum], Awaitable[mapper.Messages]]):
     server = mapper.MapAsyncServer()
 
     # Register loop-level signal handlers so we control shutdown and avoid asyncio.run

packages/pynumaflow-lite/manifests/mapstream/mapstream_cat.py

@@ -1,6 +1,7 @@
 import asyncio
 import signal
 from collections.abc import AsyncIterator
+from typing import Callable
 
 from pynumaflow_lite import mapstreamer
 from pynumaflow_lite.mapstreamer import Message
@@ -24,7 +25,7 @@ async def handler(self, keys: list[str], datum: mapstreamer.Datum) -> AsyncItera
     pass
 
 
-async def start(f: callable):
+async def start(f: Callable[[list[str], mapstreamer.Datum], AsyncIterator[Message]]):
     # Use default socket/info file locations; no explicit sock file passed
     server = mapstreamer.MapStreamAsyncServer()
 

packages/pynumaflow-lite/manifests/reduce/reduce_counter_class.py

@@ -32,7 +32,7 @@ async def handler(
     pass
 
 
-async def start(creator: type, init_args: tuple):
+async def start(creator: type[reducer.Reducer], init_args: tuple):
     sock_file = "/var/run/numaflow/reduce.sock"
     server_info_file = "/var/run/numaflow/reducer-server-info"
     server = reducer.ReduceAsyncServer(sock_file, server_info_file)

packages/pynumaflow-lite/manifests/sink/sink_log.py

@@ -2,7 +2,8 @@
 import collections
 import logging
 import signal
-from collections.abc import AsyncIterable
+from collections.abc import AsyncIterable, AsyncIterator
+from typing import Awaitable, Callable
 
 from pynumaflow_lite import sinker
 from pynumaflow_lite._sink_dtypes import Sinker
@@ -35,7 +36,7 @@ async def handler(self, datums: AsyncIterable[sinker.Datum]) -> sinker.Responses
     pass
 
 
-async def start(f: collections.abc.Callable):
+async def start(f: Callable[[AsyncIterator[sinker.Datum]], Awaitable[sinker.Responses]]):
     server = sinker.SinkAsyncServer()
 
     # Register loop-level signal handlers so we control shutdown and avoid asyncio.run

packages/pynumaflow-lite/manifests/sourcetransform/sourcetransform_event_filter.py

@@ -1,6 +1,8 @@
 import asyncio
 import signal
 from datetime import datetime, timezone
+from typing import Callable
+
 from pynumaflow_lite import sourcetransformer
 
 # Define epoch timestamps for filtering
@@ -59,7 +61,7 @@ async def handler(
     pass
 
 
-async def start(f: callable):
+async def start(f: Callable[[list[str], sourcetransformer.Datum], sourcetransformer.Messages]):
     server = sourcetransformer.SourceTransformAsyncServer()
 
     # Register loop-level signal handlers so we control shutdown and avoid asyncio.run

packages/pynumaflow-lite/pynumaflow_lite/accumulator.pyi

@@ -1,5 +1,6 @@
 from datetime import datetime
-from typing import AsyncIterator, Optional
+from typing import AsyncIterator, Optional, Callable, Awaitable
+from collections.abc import AsyncIterable
 
 class Message:
     """
@@ -74,13 +75,13 @@ class AccumulatorAsyncServer:
         info_file: str | None = "/var/run/numaflow/accumulator-server-info",
     ) -> None: ...
     async def start(
-        self, py_creator: object, init_args: object | None = None
+        self, py_creator: type[Accumulator] | Callable[[AsyncIterable[Datum]], AsyncIterator[Message]], init_args: tuple | None = None
     ) -> None:
         """
-        Start the server with the given Python class (creator).
+        Start the server with the given Python class (creator) or function.
 
         Args:
-            py_creator: The Python class to instantiate per key
+            py_creator: The Python class to instantiate per key or a function
             init_args: Optional tuple of positional arguments for class instantiation
         """
         ...

packages/pynumaflow-lite/pynumaflow_lite/batchmapper.pyi

@@ -57,7 +57,7 @@ class BatchMapAsyncServer:
             info_file: str | None = ...,
     ) -> None: ...
 
-    def start(self, py_func: Callable[..., Any]) -> Awaitable[None]: ...
+    def start(self, py_func: Callable[[AsyncIterator[Datum]], Awaitable[BatchResponses]]) -> Awaitable[None]: ...
 
     def stop(self) -> None: ...
 

packages/pynumaflow-lite/pynumaflow_lite/mapper.pyi

@@ -113,7 +113,7 @@ class MapAsyncServer:
             info_file: str | None = ...,
     ) -> None: ...
 
-    def start(self, py_func: Callable[..., Any]) -> Awaitable[None]: ...
+    def start(self, py_func: Callable[[list[str], Datum], Awaitable[Messages]]) -> Awaitable[None]: ...
 
     def stop(self) -> None: ...
 

packages/pynumaflow-lite/pynumaflow_lite/mapstreamer.pyi

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Optional, List, Dict, Callable, Awaitable, Any, AsyncIterator
+from typing import Optional, List, Dict, Callable, Awaitable, AsyncIterator
 import datetime as _dt
 
 # Re-export the Python ABC for user convenience and typing
@@ -45,7 +45,7 @@ class MapStreamAsyncServer:
             info_file: str | None = ...,
     ) -> None: ...
 
-    def start(self, py_func: Callable[..., Any]) -> Awaitable[None]: ...
+    def start(self, py_func: Callable[[list[str], Datum], AsyncIterator[Message]]) -> Awaitable[None]: ...
 
     def stop(self) -> None: ...