fix(typing): pyi, examples, and manifest

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: ...
 

packages/pynumaflow-lite/pynumaflow_lite/reducer.pyi

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 import datetime as _dt
-from typing import Optional, List, Dict, Awaitable
+from typing import Optional, List, Dict, Awaitable, Callable
+from collections.abc import AsyncIterable
 
 # Re-export the Python ABC for user convenience and typing
 from ._reduce_dtypes import Reducer as Reducer
@@ -64,7 +65,11 @@ class ReduceAsyncServer:
             info_file: str | None = ...,
     ) -> None: ...
 
-    def start(self, py_creator: type, init_args: tuple | None = ...) -> Awaitable[None]: ...
+    def start(
+        self,
+        py_creator: type[Reducer] | Callable[[list[str], AsyncIterable[Datum], Metadata], Awaitable[Messages]],
+        init_args: tuple | None = ...,
+    ) -> Awaitable[None]: ...
 
     def stop(self) -> None: ...
 

packages/pynumaflow-lite/pynumaflow_lite/reducestreamer.pyi

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 import datetime as _dt
-from typing import Optional, List, Dict, Awaitable, AsyncIterator
+from typing import Optional, List, Dict, Awaitable, AsyncIterator, Callable
+from collections.abc import AsyncIterable
 
 # Re-export the Python ABC for user convenience and typing
 from ._reducestreamer_dtypes import ReduceStreamer as ReduceStreamer
@@ -61,7 +62,7 @@ class ReduceStreamAsyncServer:
             info_file: str = ...,
     ) -> None: ...
 
-    def start(self, py_creator: type, init_args: tuple | None = ...) -> Awaitable[None]: ...
+    def start(self, py_creator: type[ReduceStreamer] | Callable[[list[str], AsyncIterable[Datum], Metadata], AsyncIterator[Message]], init_args: tuple | None = ...) -> Awaitable[None]: ...
 
     def stop(self) -> None: ...
 

packages/pynumaflow-lite/pynumaflow_lite/session_reducer.pyi

@@ -42,7 +42,7 @@ class SessionReduceAsyncServer:
             info_file: str | None = ...,
     ) -> None: ...
 
-    def start(self, py_creator: type, init_args: tuple | None = ...) -> Awaitable[None]: ...
+    def start(self, py_creator: type[SessionReducer], init_args: tuple | None = ...) -> Awaitable[None]: ...
 
     def stop(self) -> None: ...
 

packages/pynumaflow-lite/pynumaflow_lite/sinker.pyi

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

packages/pynumaflow-lite/pynumaflow_lite/sourcer.pyi

@@ -154,7 +154,7 @@ class SourceAsyncServer:
         info_file: str | None = ...,
     ) -> None: ...
 
-    def start(self, py_func: Callable[..., Any]) -> Awaitable[None]: ...
+    def start(self, py_func: Sourcer) -> Awaitable[None]: ...
 
     def stop(self) -> None: ...
 

packages/pynumaflow-lite/pynumaflow_lite/sourcetransformer.pyi

@@ -115,7 +115,7 @@ class SourceTransformAsyncServer:
             info_file: str | None = ...,
     ) -> None: ...
 
-    def start(self, py_func: Callable[..., Any]) -> Awaitable[None]: ...
+    def start(self, py_func: SourceTransformer) -> Awaitable[None]: ...
 
     def stop(self) -> None: ...
 

packages/pynumaflow-lite/tests/examples/batchmap_cat.py

@@ -1,7 +1,7 @@
 import asyncio
 import collections.abc
-
 import signal
+from typing import Awaitable, Callable
 
 from pynumaflow_lite import batchmapper
 
@@ -19,7 +19,7 @@ async def async_handler(batch: collections.abc.AsyncIterator[batchmapper.Datum])
     return responses
 
 
-async def start(f: callable):
+async def start(f: Callable[[collections.abc.AsyncIterator[batchmapper.Datum]], Awaitable[batchmapper.BatchResponses]]):
     sock_file = "/tmp/var/run/numaflow/batchmap.sock"
     server_info_file = "/tmp/var/run/numaflow/mapper-server-info"
     server = batchmapper.BatchMapAsyncServer(sock_file, server_info_file)

packages/pynumaflow-lite/tests/examples/batchmap_cat_class.py

@@ -1,6 +1,7 @@
 import asyncio
 import signal
 from collections.abc import AsyncIterator
+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: AsyncIterator[batchmapper.Datum]) -> batchmapper.
     pass
 
 
-async def start(f: callable):
+async def start(f: Callable[[AsyncIterator[batchmapper.Datum]], Awaitable[batchmapper.BatchResponses]]):
     sock_file = "/tmp/var/run/numaflow/batchmap.sock"
     server_info_file = "/tmp/var/run/numaflow/mapper-server-info"
     server = batchmapper.BatchMapAsyncServer(sock_file, server_info_file)

packages/pynumaflow-lite/tests/examples/map_cat.py

@@ -1,5 +1,6 @@
 import asyncio
 import signal
+from typing import Awaitable, Callable
 
 from pynumaflow_lite import mapper
 
@@ -37,7 +38,7 @@ async def async_handler(
     return messages
 
 
-async def start(f: callable):
+async def start(f: Callable[[list[str], mapper.Datum], Awaitable[mapper.Messages]]):
     sock_file = "/tmp/var/run/numaflow/map.sock"
     server_info_file = "/tmp/var/run/numaflow/mapper-server-info"
     server = mapper.MapAsyncServer(sock_file, server_info_file)

packages/pynumaflow-lite/tests/examples/map_cat_class.py

@@ -1,5 +1,7 @@
 import asyncio
 import signal
+from typing import Awaitable, Callable
+
 from pynumaflow_lite import mapper
 
 
@@ -46,7 +48,7 @@ async def handler(
     pass
 
 
-async def start(f: callable):
+async def start(f: Callable[[list[str], mapper.Datum], Awaitable[mapper.Messages]]):
     sock_file = "/tmp/var/run/numaflow/map.sock"
     server_info_file = "/tmp/var/run/numaflow/mapper-server-info"
     server = mapper.MapAsyncServer(sock_file, server_info_file)

packages/pynumaflow-lite/tests/examples/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
@@ -19,7 +20,7 @@ async def async_handler(keys: list[str], datum: mapstreamer.Datum) -> AsyncItera
         yield Message(s.encode(), keys)
 
 
-async def start(f: callable):
+async def start(f: Callable[[list[str], mapstreamer.Datum], AsyncIterator[Message]]):
     sock_file = "/tmp/var/run/numaflow/mapstream.sock"
     server_info_file = "/tmp/var/run/numaflow/mapper-server-info"
     server = mapstreamer.MapStreamAsyncServer(sock_file, server_info_file)

packages/pynumaflow-lite/tests/examples/mapstream_cat_class.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]]):
     sock_file = "/tmp/var/run/numaflow/mapstream.sock"
     server_info_file = "/tmp/var/run/numaflow/mapper-server-info"
     server = mapstreamer.MapStreamAsyncServer(sock_file, server_info_file)

packages/pynumaflow-lite/tests/examples/reduce_counter_func.py

@@ -1,6 +1,7 @@
 import asyncio
 import signal
 from collections.abc import AsyncIterable
+from typing import Awaitable, Callable
 
 from pynumaflow_lite import reducer
 
@@ -29,7 +30,9 @@ async def reduce_handler(
     pass
 
 
-async def start(handler: callable):
+async def start(
+    handler: Callable[[list[str], AsyncIterable[reducer.Datum], reducer.Metadata], Awaitable[reducer.Messages]]
+):
     sock_file = "/tmp/var/run/numaflow/reduce.sock"
     server_info_file = "/tmp/var/run/numaflow/reducer-server-info"
     server = reducer.ReduceAsyncServer(sock_file, server_info_file)

packages/pynumaflow-lite/tests/examples/sink_log.py

@@ -2,6 +2,7 @@
 import collections.abc
 import logging
 import signal
+from typing import Awaitable, Callable
 
 from pynumaflow_lite import sinker
 
@@ -39,7 +40,7 @@ async def async_handler(datums: collections.abc.AsyncIterator[sinker.Datum]) ->
     return responses
 
 
-async def start(f: callable):
+async def start(f: Callable[[collections.abc.AsyncIterator[sinker.Datum]], Awaitable[sinker.Responses]]):
     sock_file = "/tmp/var/run/numaflow/sink.sock"
     server_info_file = "/tmp/var/run/numaflow/sinker-server-info"
     server = sinker.SinkAsyncServer(sock_file, server_info_file)

packages/pynumaflow-lite/tests/examples/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
@@ -91,7 +93,7 @@ async def handler(
     pass
 
 
-async def start(f: callable):
+async def start(f: Callable[[list[str], sourcetransformer.Datum], sourcetransformer.Messages]):
     sock_file = "/tmp/var/run/numaflow/sourcetransform.sock"
     server_info_file = "/tmp/var/run/numaflow/sourcetransformer-server-info"
     server = sourcetransformer.SourceTransformAsyncServer(sock_file, server_info_file)