type + misc

src/wf_authoring/nodes/callables.py

@@ -9,48 +9,45 @@ from wf_core import RuntimeContext
 
 from .result import NodeReturn
 
-InputT = TypeVar("InputT", bound=BaseModel)
-OutputT = TypeVar("OutputT", bound=BaseModel)
-
-InputT_contra = TypeVar("InputT_contra", bound=BaseModel, contravariant=True)
-OutputT_co = TypeVar("OutputT_co", bound=BaseModel, covariant=True)
+InputT = TypeVar("InputT", bound=BaseModel, infer_variance=True)
+OutputT = TypeVar("OutputT", bound=BaseModel, infer_variance=True)
 
 
-class ContextNodeCallable(Protocol[InputT_contra, OutputT_co]):
+class ContextNodeCallable(Protocol[InputT, OutputT]):
     def __call__(
         self,
-        payload: InputT_contra,
+        payload: InputT,
         /,
         ctx: RuntimeContext,
-    ) -> NodeReturn[OutputT_co] | OutputT_co | None: ...
+    ) -> NodeReturn[OutputT] | OutputT | None: ...
 
 
-class PlainNodeCallable(Protocol[InputT_contra, OutputT_co]):
+class PlainNodeCallable(Protocol[InputT, OutputT]):
     def __call__(
         self,
-        payload: InputT_contra,
+        payload: InputT,
         /,
-    ) -> NodeReturn[OutputT_co] | OutputT_co | None: ...
+    ) -> NodeReturn[OutputT] | OutputT | None: ...
 
 
 NodeCallable = ContextNodeCallable[InputT, OutputT] | PlainNodeCallable[InputT, OutputT]
 
 
-class AsyncContextNodeCallable(Protocol[InputT_contra, OutputT_co]):
+class AsyncContextNodeCallable(Protocol[InputT, OutputT]):
     def __call__(
         self,
-        payload: InputT_contra,
+        payload: InputT,
         /,
         ctx: RuntimeContext,
-    ) -> Awaitable[NodeReturn[OutputT_co] | OutputT_co | None]: ...
+    ) -> Awaitable[NodeReturn[OutputT] | OutputT | None]: ...
 
 
-class AsyncPlainNodeCallable(Protocol[InputT_contra, OutputT_co]):
+class AsyncPlainNodeCallable(Protocol[InputT, OutputT]):
     def __call__(
         self,
-        payload: InputT_contra,
+        payload: InputT,
         /,
-    ) -> Awaitable[NodeReturn[OutputT_co] | OutputT_co | None]: ...
+    ) -> Awaitable[NodeReturn[OutputT] | OutputT | None]: ...
 
 
 AsyncNodeCallable = (

src/wf_authoring/reducers/callables.py

@@ -4,7 +4,7 @@ from typing import Any, Protocol, TypeVar
 
 from pydantic import BaseModel
 
-ConfigT = TypeVar("ConfigT", bound=BaseModel)
+ConfigT = TypeVar("ConfigT", bound=BaseModel, infer_variance=True)
 
 
 class PlainReducerCallable(Protocol):

src/wf_authoring/reducers/decorator.py

@@ -2,17 +2,13 @@ from __future__ import annotations
 
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass
-from typing import Any, TypeVar, cast, overload
-
-from pydantic import BaseModel
+from typing import Any, cast, overload
 
 from wf_core import ReducerSpec, SiblingWritePolicy
 from wf_core.runtime.ops.merges import ReducerDefinition
 
 from .callables import ConfigReducerCallable, ConfigT, PlainReducerCallable
 
-PlainFnT = TypeVar("PlainFnT", bound=PlainReducerCallable)
-
 
 @dataclass(frozen=True, slots=True)
 class AuthoredReducer:
@@ -21,9 +17,10 @@ class AuthoredReducer:
     definition: ReducerDefinition
 
 
+# @reducer no parentheses / manual overload
 @overload
 def reducer(
-    fn: PlainFnT,
+    fn: PlainReducerCallable,
     /,
     *,
     name: str | None = None,
@@ -32,15 +29,29 @@ def reducer(
 ) -> AuthoredReducer: ...
 
 
+@overload
+def reducer(
+    fn: ConfigReducerCallable[ConfigT],
+    /,
+    *,
+    name: str | None = None,
+    config_model: type[ConfigT],
+    description: str | None = None,
+    sibling_write_policy: SiblingWritePolicy = SiblingWritePolicy.MERGEABLE,
+) -> AuthoredReducer: ...
+
+
+# configless overload
 @overload
 def reducer(
     *,
     name: str,
     description: str | None = None,
     sibling_write_policy: SiblingWritePolicy = SiblingWritePolicy.MERGEABLE,
-) -> Callable[[Callable[..., Any]], AuthoredReducer]: ...
+) -> Callable[[PlainReducerCallable], AuthoredReducer]: ...
 
 
+# config overload
 @overload
 def reducer(
     *,
@@ -48,24 +59,38 @@ def reducer(
     config_model: type[ConfigT],
     description: str | None = None,
     sibling_write_policy: SiblingWritePolicy = SiblingWritePolicy.MERGEABLE,
-) -> Callable[[Callable[..., Any]], AuthoredReducer]: ...
+) -> Callable[[ConfigReducerCallable[ConfigT]], AuthoredReducer]: ...
 
 
 def reducer(
-    fn: Callable[..., Any] | None = None,
+    fn: PlainReducerCallable | ConfigReducerCallable[ConfigT] | None = None,
     /,
     *,
     name: str | None = None,
-    config_model: type[BaseModel] | None = None,
+    config_model: type[ConfigT] | None = None,
     description: str | None = None,
     sibling_write_policy: SiblingWritePolicy = SiblingWritePolicy.MERGEABLE,
-) -> AuthoredReducer | Callable[[Callable[..., Any]], AuthoredReducer]:
-    """Wrap a Python reducer function as a runtime reducer definition."""
+) -> (
+    AuthoredReducer
+    | Callable[[PlainReducerCallable], AuthoredReducer]
+    | Callable[[ConfigReducerCallable[ConfigT]], AuthoredReducer]
+):
+    """Wrap a Python reducer function as a runtime reducer definition.
 
-    def decorate(raw: Callable[..., Any]) -> AuthoredReducer:
-        reducer_name = name or raw.__name__
-        reducer_description = description or raw.__doc__
-        if config_model is None:
+    Assume wrapped function will not approve of the third `config` argument if `config_model` is not provided, and vice versa
+    """
+    # how do i even do this...
+    # if @overload return is (PlainReducerCallable) -> Authored. (or even another is (ConfigReducerCallable[ConfigT]) -> Authored). then the base signature must not be (Plain | Config) -> Authored???
+    # why is that, basedpyright?
+
+    if config_model is None:
+
+        def decorate_plain(
+            raw: PlainReducerCallable,
+        ) -> AuthoredReducer:
+            reducer_name = name or raw.__name__
+            reducer_description = description or raw.__doc__
+            # if no config model is provided, we assume it's a plain reducer and just wrap it directly
             return AuthoredReducer(
                 ReducerDefinition(
                     spec=ReducerSpec(
@@ -77,36 +102,49 @@ def reducer(
                 )
             )
 
-        model_type = config_model
+        if fn is not None:
+            return decorate_plain(cast(PlainReducerCallable, fn))
+        return decorate_plain
 
-        def runtime_fn(
-            current: Any,
-            incoming: Any,
-            config: Mapping[str, Any],
-        ) -> Any:
-            parsed = model_type.model_validate(config)
-            return cast(ConfigReducerCallable[BaseModel], cast(object, raw))(
-                current,
-                incoming,
-                parsed,
+    else:
+
+        def decorate_config(
+            raw: ConfigReducerCallable[ConfigT],
+        ) -> AuthoredReducer:
+            reducer_name = name or raw.__name__
+            reducer_description = description or raw.__doc__
+
+            model_type = config_model
+
+            # if a config model is provided, we need to parse the config before calling the reducer function
+            def runtime_fn(
+                current: Any,
+                incoming: Any,
+                config: Mapping[str, Any],
+            ) -> Any:
+                parsed = model_type.model_validate(config)
+                return raw(
+                    current,
+                    incoming,
+                    parsed,
+                )
+
+            return AuthoredReducer(
+                ReducerDefinition(
+                    spec=ReducerSpec(
+                        name=reducer_name,
+                        description=_clean_doc(reducer_description),
+                        config_schema=config_model.model_json_schema(),
+                        sibling_write_policy=sibling_write_policy,
+                    ),
+                    fn=runtime_fn,
+                    accepts_config=True,
+                )
             )
 
-        return AuthoredReducer(
-            ReducerDefinition(
-                spec=ReducerSpec(
-                    name=reducer_name,
-                    description=_clean_doc(reducer_description),
-                    config_schema=config_model.model_json_schema(),
-                    sibling_write_policy=sibling_write_policy,
-                ),
-                fn=runtime_fn,
-                accepts_config=True,
-            )
-        )
-
-    if fn is not None:
-        return decorate(fn)
-    return decorate
+        if fn is not None:
+            return decorate_config(cast(ConfigReducerCallable[ConfigT], fn))
+        return decorate_config
 
 
 def _clean_doc(doc: str | None) -> str | None: