Typing

Author: codeSamuraii

examples/async_execution.py

@@ -14,7 +14,7 @@
 
 pyfuse.connect("local://localhost:9748")
 
-async def add(a: int, b: int) -> int:
+async def add(a: float, b: float) -> float:
     return a + b
 
 @trace
@@ -39,12 +39,12 @@ async def main() -> None:
     print(f"map results = {results}")  # [5.0, 13.0, 17.0]
 
     # 4. asyncio.gather -- submit multiple tasks, await concurrently
-    results = await asyncio.gather(
+    gathered = await asyncio.gather(
         hypotenuse.run(3.0, 4.0),
         hypotenuse.run(5.0, 12.0),
         hypotenuse.run(8.0, 15.0),
     )
-    print(f"gather: {results}")  # [5.0, 13.0, 17.0]
+    print(f"gather: {gathered}")  # [5.0, 13.0, 17.0]
 
 
 asyncio.run(main())

examples/remote_execution.py

@@ -1,12 +1,19 @@
 import asyncio
 import math
+from typing import overload
 
 import pyfuse
 from pyfuse import trace
 
 pyfuse.connect("local://localhost:9748")
 
-def add(a: int, b: int) -> int:
+@overload
+def add(a: int, b: int) -> int: ...
+
+@overload
+def add(a: float, b: float) -> float: ...
+
+def add(a: int | float, b: int | float) -> int | float:
     return a + b
 
 @trace

pyfuse/graph/decorator.py

@@ -1,27 +1,29 @@
 import logging
 from collections.abc import Callable
-from typing import TypeVar, overload
+from typing import ParamSpec, TypeVar, overload
 
 from pyfuse.graph.graph import Graph
+from pyfuse.typing import TraceDecorator, TracedFunction
 
 logger = logging.getLogger(__name__)
 
-_F = TypeVar("_F", bound=Callable[..., object])
+_R = TypeVar("_R")
+_P = ParamSpec("_P")
 
 
 @overload
-def trace(func: _F) -> _F: ...
+def trace(func: Callable[_P, _R]) -> TracedFunction[_P, _R]: ...
 @overload
-def trace(*, timeout: float | None = ..., retries: int = ..., retry_delay: float = ...) -> Callable[[_F], _F]: ...
+def trace(*, timeout: float | None = ..., retries: int = ..., retry_delay: float = ...) -> TraceDecorator: ...
 
 
 def trace(
-    func: _F | None = None,
+    func: Callable[..., object] | None = None,
     *,
     timeout: float | None = None,
     retries: int = 0,
     retry_delay: float = 1.0,
-) -> _F | Callable[[_F], _F]:
+) -> object:
     """Enable a function for serialization and remote execution.
 
     The decorated function works normally when called directly.
@@ -38,19 +40,19 @@ def flaky(x): ...
     if func is not None:
         return _apply_trace(func, timeout=timeout, retries=retries, retry_delay=retry_delay)
 
-    def decorator(f: _F) -> _F:
+    def decorator(f: Callable[_P, _R]) -> object:
         return _apply_trace(f, timeout=timeout, retries=retries, retry_delay=retry_delay)
 
     return decorator
 
 
 def _apply_trace(
-    func: _F,
+    func: Callable[_P, _R],
     *,
     timeout: float | None = None,
     retries: int = 0,
     retry_delay: float = 1.0,
-) -> _F:
+) -> TracedFunction[_P, _R]:
     logger.debug("@trace applied to %s", func.__qualname__)
     graph = Graph.default()
     graph.register(func)

pyfuse/graph/tracing.py

@@ -9,13 +9,16 @@
 import threading
 from collections.abc import AsyncGenerator, Callable, Generator
 from pathlib import Path
-from typing import Any, TypeVar
+from typing import Any, ParamSpec, TypeVar, cast
 
+from pyfuse.typing import TracedFunction
 from pyfuse.worker.backends.base import Backend
 
 logger = logging.getLogger(__name__)
 
 _F = TypeVar("_F", bound=Callable[..., object])
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
 
 _BUILTIN_NAMES = set(dir(builtins))
 
@@ -164,7 +167,7 @@ def _record_edge(self, stack: list[str], qualified_name: str) -> None:
             with self._lock:
                 self._runtime_deps.setdefault(stack[-1], set()).add(qualified_name)
 
-    def create_wrapper(self, func: _F) -> _F:
+    def create_wrapper(self, func: Callable[_P, _R]) -> TracedFunction[_P, _R]:
         """Wrap func to record runtime caller-callee edges.
 
         The wrapper preserves the original function signature via
@@ -182,7 +185,7 @@ def create_wrapper(self, func: _F) -> _F:
             wrapper = self._wrap_generator(func, qualified_name)
         else:
             wrapper = self._wrap_sync(func, qualified_name)
-        return wrapper  # type: ignore[no-any-return]
+        return cast(TracedFunction[_P, _R], wrapper)
 
     def _wrap_async_generator(self, func: Any, qualified_name: str) -> Any:
         logger.debug("Creating async generator wrapper for %s", qualified_name)

pyfuse/typing.py

@@ -0,0 +1,24 @@
+from collections.abc import Callable
+from typing import Any, ParamSpec, Protocol, TypeVar
+
+from pyfuse.worker.result import Result
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+class TracedFunction(Protocol[P, R]):
+    __pyfuse_traced__: bool
+    __wrapped__: Callable[P, R]
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R: ...
+
+    async def start(self, *args: P.args, **kwargs: P.kwargs) -> Result: ...
+
+    async def run(self, *args: P.args, **kwargs: P.kwargs) -> Any: ...
+
+    async def map(self, args_list: list[tuple[Any, ...]], **kwargs: Any) -> list[Any]: ...
+
+
+class TraceDecorator(Protocol):
+    def __call__(self, func: Callable[P, R]) -> TracedFunction[P, R]: ...

pyfuse/worker/backends/rabbitmq.py

@@ -17,7 +17,7 @@
 from typing import Any
 
 try:
-    import aio_pika
+    import aio_pika  # type: ignore[import-not-found]
 except ImportError:
     raise ImportError(
         "aio-pika package is required for RabbitMQBackend. "

tests/test_packaging.py

@@ -15,8 +15,6 @@
 
 from pyfuse.core.version import _FALLBACK_VERSION, _VERSION
 
-pytestmark = pytest.mark.slow
-
 
 def _project_root() -> Path:
     """Return the repository root (contains pyproject.toml)."""

tests/test_venv.py

@@ -25,8 +25,6 @@
     temp_venv,
 )
 
-pytestmark = pytest.mark.slow
-
 
 class TestFindProjectRoot:
     def test_finds_root(self) -> None: