feat: refactor type detection in decorator

Author: brian14708

src/duron/_core/context.py

@@ -18,6 +18,7 @@
 from duron._core.signal import create_signal
 from duron._core.stream import create_stream, run_stream
 from duron._decorator.op import CheckpointOp, Op
+from duron.typing import inspect_function
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Coroutine
@@ -41,11 +42,11 @@
 
 @final
 class Context:
-    __slots__ = ("_loop", "_task", "_token")
+    __slots__ = ("_fn", "_loop", "_token")
 
     def __init__(self, task: Fn[..., object], loop: EventLoop) -> None:
         self._loop: EventLoop = loop
-        self._task = task
+        self._fn = task
         self._token: Token[Context | None] | None = None
 
     def __enter__(self) -> Context:
@@ -109,12 +110,9 @@ async def run(
                 return await stream
 
         if isinstance(fn, Op):
-            if fn.return_type:
-                return_type = fn.return_type
-            else:
-                return_type = self._task.codec.inspect_function(fn.fn).return_type
+            return_type = fn.return_type
         else:
-            return_type = self._task.codec.inspect_function(fn).return_type
+            return_type = inspect_function(fn).return_type
 
         metadata = options.metadata if options else None
         op = create_op(

src/duron/_core/invoke.py

@@ -31,7 +31,7 @@
 from duron._loop import EventLoop, create_loop
 from duron.codec import Codec, JSONValue
 from duron.log import is_entry
-from duron.typing import unspecified
+from duron.typing import inspect_function, unspecified
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Coroutine
@@ -44,7 +44,7 @@
     from duron._core.stream import Stream
     from duron._decorator.fn import Fn
     from duron._loop import OpFuture, WaitSet
-    from duron.codec import Codec, FunctionType
+    from duron.codec import Codec
     from duron.log import (
         BarrierEntry,
         Entry,
@@ -56,7 +56,7 @@
         StreamCreateEntry,
         StreamEmitEntry,
     )
-    from duron.typing import TypeHint
+    from duron.typing import FunctionType, TypeHint
 
 
 _T_co = TypeVar("_T_co", covariant=True)
@@ -100,7 +100,7 @@ def get_init() -> InitParams:
             }
 
         codec = self._fn.codec
-        type_info = codec.inspect_function(self._fn.fn)
+        type_info = inspect_function(self._fn.fn)
         prelude = _invoke_prelude(self._fn, type_info, get_init)
         self._run = _InvokeRun(
             prelude,
@@ -115,7 +115,7 @@ def cb() -> InitParams:
             msg = "not started"
             raise RuntimeError(msg)
 
-        type_info = self._fn.codec.inspect_function(self._fn.fn)
+        type_info = inspect_function(self._fn.fn)
         prelude = _invoke_prelude(self._fn, type_info, cb)
         self._run = _InvokeRun(
             prelude,

src/duron/_decorator/op.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections.abc import AsyncGenerator
 from dataclasses import dataclass
 from typing import (
     TYPE_CHECKING,
@@ -9,13 +10,15 @@
     Literal,
     ParamSpec,
     TypeVar,
+    get_args,
+    get_origin,
     overload,
 )
 
-from duron.typing import unspecified
+from duron.typing import inspect_function, unspecified
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Callable, Coroutine
+    from collections.abc import Callable, Coroutine
 
     from duron.typing import TypeHint
 
@@ -32,7 +35,7 @@ class CheckpointOp(Generic[_P, _S, _T]):
     initial: Callable[[], _S]
     reducer: Callable[[_S, _T], _S]
     action_type: TypeHint[_T]
-    state_type: TypeHint[_S]
+    return_type: TypeHint[_S]
 
     def __call__(
         self,
@@ -108,20 +111,33 @@ def op(
         def decorate_ckpt(
             fn: Callable[Concatenate[_S, _P], AsyncGenerator[_T, _S]],
         ) -> CheckpointOp[_P, _S, _T]:
+            action_type_local = action_type
+            return_type_local = return_type
+            if not action_type_local or not return_type_local:
+                ret = inspect_function(fn).return_type
+                if get_origin(ret) is AsyncGenerator:
+                    yield_, send = get_args(ret)
+                    if not action_type_local:
+                        action_type_local = yield_
+                    if not return_type_local:
+                        return_type_local = send
             return CheckpointOp(
                 fn=fn,
                 initial=initial,
                 reducer=reducer,
-                action_type=action_type,
-                state_type=return_type,
+                action_type=action_type_local,
+                return_type=return_type_local,
             )
 
         return decorate_ckpt
 
     def decorate(
         fn: Callable[_P, Coroutine[Any, Any, _T_co]],
     ) -> Op[_P, _T_co]:
-        return Op(fn=fn, return_type=return_type)
+        return Op(
+            fn=fn,
+            return_type=return_type or inspect_function(fn).return_type,
+        )
 
     if fn is not None:
         return decorate(fn)

src/duron/codec.py

@@ -1,15 +1,10 @@
 from __future__ import annotations
 
-import inspect
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
 from typing import TYPE_CHECKING, TypeGuard, cast, final
 from typing_extensions import TypeAliasType, override
 
-from duron.typing import unspecified
-
 if TYPE_CHECKING:
-    from collections.abc import Callable
     from typing_extensions import Any
 
     from duron.typing import TypeHint
@@ -37,13 +32,6 @@ def is_json_value(x: object) -> TypeGuard[JSONValue]:
     return False
 
 
-@dataclass(slots=True)
-class FunctionType:
-    return_type: TypeHint[Any]
-    parameters: list[str]
-    parameter_types: dict[str, TypeHint[Any]]
-
-
 class Codec(ABC):
     __slots__: tuple[str, ...] = ()
 
@@ -58,40 +46,6 @@ def decode_json(
         /,
     ) -> object: ...
 
-    def inspect_function(  # noqa: PLR6301
-        self,
-        fn: Callable[..., object],
-    ) -> FunctionType:
-        try:
-            sig = inspect.signature(fn, eval_str=True)
-        except NameError:
-            sig = inspect.signature(fn)
-        return_type = (
-            sig.return_annotation
-            if sig.return_annotation != inspect.Parameter.empty
-            else unspecified
-        )
-
-        parameter_names: list[str] = []
-        parameter_types: dict[str, TypeHint[Any]] = {}
-        for k, p in sig.parameters.items():
-            if p.kind in {
-                inspect.Parameter.VAR_POSITIONAL,
-                inspect.Parameter.VAR_KEYWORD,
-            }:
-                continue
-
-            parameter_names.append(k)
-            parameter_types[p.name] = (
-                p.annotation if p.annotation != inspect.Parameter.empty else unspecified
-            )
-
-        return FunctionType(
-            return_type=return_type,
-            parameters=parameter_names,
-            parameter_types=parameter_types,
-        )
-
 
 @final
 class DefaultCodec(Codec):

src/duron/typing/__init__.py

@@ -0,0 +1,4 @@
+from .hint import TypeHint as TypeHint
+from .hint import unspecified as unspecified
+from .inspect import FunctionType as FunctionType
+from .inspect import inspect_function as inspect_function

src/duron/typing.py src/duron/typing/hint.pysrc/duron/typing.py renamed to src/duron/typing/hint.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from types import UnionType
 from typing_extensions import TypeAliasType, TypeVar
 
@@ -23,5 +25,3 @@ def __bool__(self) -> bool:
     from typing_extensions import TypeForm
 
     TypeHint = TypeAliasType("TypeHint", TypeForm[_T] | _Unspecified, type_params=(_T,))
-
-__all__ = ["TypeHint", "unspecified"]

src/duron/typing/inspect.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from duron.typing.hint import unspecified
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from duron.typing.hint import TypeHint
+
+
+@dataclass(slots=True)
+class FunctionType:
+    return_type: TypeHint[Any]
+    parameters: list[str]
+    parameter_types: dict[str, TypeHint[Any]]
+
+
+def inspect_function(
+    fn: Callable[..., object],
+) -> FunctionType:
+    try:
+        sig = inspect.signature(fn, eval_str=True)
+    except NameError:
+        sig = inspect.signature(fn)
+    return_type = (
+        sig.return_annotation
+        if sig.return_annotation != inspect.Parameter.empty
+        else unspecified
+    )
+
+    parameter_names: list[str] = []
+    parameter_types: dict[str, TypeHint[Any]] = {}
+    for k, p in sig.parameters.items():
+        if p.kind in {
+            inspect.Parameter.VAR_POSITIONAL,
+            inspect.Parameter.VAR_KEYWORD,
+        }:
+            continue
+
+        parameter_names.append(k)
+        parameter_types[p.name] = (
+            p.annotation if p.annotation != inspect.Parameter.empty else unspecified
+        )
+
+    return FunctionType(
+        return_type=return_type,
+        parameters=parameter_names,
+        parameter_types=parameter_types,
+    )

tests/test_codec.py tests/test_typing.pytests/test_codec.py renamed to tests/test_typing.py

@@ -1,21 +1,15 @@
 from __future__ import annotations
 
 from collections.abc import AsyncGenerator
-from typing import TYPE_CHECKING
 
-from duron.codec import DefaultCodec, FunctionType
-
-if TYPE_CHECKING:
-    from duron.codec import Codec
-
-codec: Codec = DefaultCodec()
+from duron.typing import FunctionType, inspect_function
 
 
 def test_no_parameters() -> None:
     def simple_func() -> int:
         return 42
 
-    result = codec.inspect_function(simple_func)
+    result = inspect_function(simple_func)
 
     assert isinstance(result, FunctionType)
     assert result.parameters == []
@@ -27,7 +21,7 @@ def test_no_parameters_no_return_type() -> None:
     def simple_func() -> None:
         pass
 
-    result = codec.inspect_function(simple_func)
+    result = inspect_function(simple_func)
 
     assert isinstance(result, FunctionType)
     assert result.parameters == []
@@ -39,7 +33,7 @@ def test_type_annotated_parameters() -> None:
     def typed_func(_x: int, _y: str, _z: float) -> bool:
         return True
 
-    result = codec.inspect_function(typed_func)
+    result = inspect_function(typed_func)
 
     assert result.parameters == ["_x", "_y", "_z"]
     assert result.parameter_types == {"_x": int, "_y": str, "_z": float}
@@ -50,7 +44,7 @@ def test_return_type_annotation() -> None:
     def func_with_return() -> dict[str, int]:
         return {}
 
-    result = codec.inspect_function(func_with_return)
+    result = inspect_function(func_with_return)
 
     assert result.parameters == []
     assert result.parameter_types == {}
@@ -61,7 +55,7 @@ def test_complex_type_annotations() -> None:
     def complex_func(_x: int | str, _y: list[dict[str, int]]) -> tuple[int, ...]:
         return (1, 2, 3)
 
-    result = codec.inspect_function(complex_func)
+    result = inspect_function(complex_func)
 
     assert result.parameters == ["_x", "_y"]
     assert result.parameter_types == {"_x": int | str, "_y": list[dict[str, int]]}
@@ -72,7 +66,7 @@ def test_async_function() -> None:
     async def async_func(_a: int, _b: str) -> bool:  # noqa: RUF029
         return True
 
-    result = codec.inspect_function(async_func)
+    result = inspect_function(async_func)
 
     assert result.parameters == ["_a", "_b"]
     assert result.parameter_types == {"_a": int, "_b": str}
@@ -83,7 +77,7 @@ def test_varargs_and_kwargs() -> None:
     def func_with_varargs(_a: int, *_args: str, **_kwargs: bool) -> None:
         pass
 
-    result = codec.inspect_function(func_with_varargs)
+    result = inspect_function(func_with_varargs)
 
     assert result.parameters == ["_a"]
     assert result.parameter_types == {"_a": int}
@@ -94,7 +88,7 @@ def test_positional_only_and_keyword_only() -> None:
     def func_with_special_args(_a: int, /, _b: str, *, _c: float) -> bool:
         return True
 
-    result = codec.inspect_function(func_with_special_args)
+    result = inspect_function(func_with_special_args)
 
     assert result.parameters == ["_a", "_b", "_c"]
     assert result.parameter_types == {"_a": int, "_b": str, "_c": float}
@@ -105,7 +99,7 @@ def test_iterator() -> None:
     async def generator() -> AsyncGenerator[int]:  # noqa: RUF029
         yield 1
 
-    result = codec.inspect_function(generator)
+    result = inspect_function(generator)
 
     assert result.parameters == []
     assert result.parameter_types == {}