Simplify Flow.model code, update docs

Signed-off-by: Nijat Khanbabayev <nijat.khanbabayev@cubistsystematic.com>

Author: Nijat Khanbabayev

ccflow/callable.py

@@ -16,7 +16,24 @@
 import logging
 from functools import lru_cache, wraps
 from inspect import Signature, isclass, signature
-from typing import TYPE_CHECKING, Any, Callable, ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union, cast, get_args, get_origin
+from typing import (
+    TYPE_CHECKING,
+    Annotated,
+    Any,
+    Callable,
+    ClassVar,
+    Dict,
+    Generic,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+    get_args,
+    get_origin,
+)
 
 from pydantic import BaseModel as PydanticBaseModel, ConfigDict, Field, InstanceOf, PrivateAttr, TypeAdapter, field_validator, model_validator
 from typing_extensions import override
@@ -65,11 +82,11 @@ def _cached_signature(fn):
     return signature(fn)
 
 
-def _callable_qualname(fn: Callable[..., Any]) -> str:
-    return getattr(fn, "__qualname__", type(fn).__qualname__)
-
-
 def _declared_type_matches(actual: Any, expected: Any) -> bool:
+    while get_origin(actual) is Annotated:
+        actual = get_args(actual)[0]
+    while get_origin(expected) is Annotated:
+        expected = get_args(expected)[0]
     if isinstance(expected, TypeVar):
         return True
     if get_origin(expected) is Union:
@@ -293,25 +310,13 @@ def wrapper(model, context=Signature.empty, *, _options: Optional[FlowOptions] =
             if not isinstance(model, CallableModel):
                 raise TypeError(f"Can only decorate methods on CallableModels (not {type(model)}) with the flow decorator.")
 
-            # Check if this is an auto_context decorated method
-            has_auto_context = hasattr(fn, "__auto_context__")
-            if has_auto_context:
-                method_context_type = fn.__auto_context__
-            else:
-                method_context_type = model.context_type
-
-            # Validate context type (skip for auto contexts which are always valid ContextBase subclasses)
-            if not has_auto_context:
-                if (not isclass(model.context_type) or not issubclass(model.context_type, ContextBase)) and not (
-                    get_origin(model.context_type) is Union and type(None) in get_args(model.context_type)
-                ):
-                    raise TypeError(f"Context type {model.context_type} must be a subclass of ContextBase")
-
-            # Validate result type - use __result_type__ for auto contexts if available
-            if has_auto_context and hasattr(fn, "__result_type__"):
-                method_result_type = fn.__result_type__
-            else:
-                method_result_type = model.result_type
+            method_context_type = getattr(fn, "__auto_context__", model.context_type)
+            method_result_type = getattr(fn, "__result_type__", model.result_type)
+
+            if (not isclass(method_context_type) or not issubclass(method_context_type, ContextBase)) and not (
+                get_origin(method_context_type) is Union and type(None) in get_args(method_context_type)
+            ):
+                raise TypeError(f"Context type {method_context_type} must be a subclass of ContextBase")
             if (not isclass(method_result_type) or not issubclass(method_result_type, ResultBase)) and not (
                 get_origin(method_result_type) is Union and all(isclass(t) and issubclass(t, ResultBase) for t in get_args(method_result_type))
             ):
@@ -332,12 +337,13 @@ def wrapper(model, context=Signature.empty, *, _options: Optional[FlowOptions] =
                 raise TypeError(f"{fn.__name__}() was passed a context and got an unexpected keyword argument '{next(iter(kwargs.keys()))}'")
 
             # Type coercion on input. We do this here (rather than relying on ModelEvaluationContext) as it produces a nicer traceback/error message
-            if not isinstance(context, method_context_type):
-                if get_origin(method_context_type) is Union and type(None) in get_args(method_context_type):
-                    coerce_context_type = [t for t in get_args(method_context_type) if t is not type(None)][0]
-                else:
-                    coerce_context_type = method_context_type
-                context = coerce_context_type.model_validate(context)
+            if get_origin(method_context_type) is Union and type(None) in get_args(method_context_type):
+                if context is not None:
+                    method_context_type = [t for t in get_args(method_context_type) if t is not type(None)][0]
+                    if not isinstance(context, method_context_type):
+                        context = method_context_type.model_validate(context)
+            elif not isinstance(context, method_context_type):
+                context = method_context_type.model_validate(context)
 
             if fn != getattr(model.__class__, fn.__name__).__wrapped__:
                 # This happens when super().__call__ is used when implementing a CallableModel that derives from another one.
@@ -356,7 +362,6 @@ def wrapper(model, context=Signature.empty, *, _options: Optional[FlowOptions] =
         wrap_any.get_options = self.get_options
         wrap_any.get_evaluation_context = get_evaluation_context
 
-        # Preserve auto context attributes for introspection
         if hasattr(fn, "__auto_context__"):
             wrap_any.__auto_context__ = fn.__auto_context__
         if hasattr(fn, "__result_type__"):
@@ -476,19 +481,12 @@ def __call__(self, *, date: date, extra: int = 0) -> MyResult:
             # with infrastructure expecting DateContext instances.
 
         """
-        # Extract auto_context option (not part of FlowOptions)
-        # Can be: False, True, or a ContextBase subclass
         auto_context = kwargs.pop("auto_context", False)
-
-        # Determine if auto_context is enabled and extract parent class if provided
         if auto_context is False:
-            auto_context_enabled = False
             context_parent = None
         elif auto_context is True:
-            auto_context_enabled = True
-            context_parent = None
+            context_parent = ContextBase
         elif isclass(auto_context) and issubclass(auto_context, ContextBase):
-            auto_context_enabled = True
             context_parent = auto_context
         else:
             raise TypeError(f"auto_context must be False, True, or a ContextBase subclass, got {auto_context!r}")
@@ -501,7 +499,7 @@ def __call__(self, *, date: date, extra: int = 0) -> MyResult:
         else:
             # Arguments to decorator, this is just returning the decorator
             # Note that the code below is executed only once
-            if auto_context_enabled:
+            if context_parent is not None:
                 # Return a decorator that first applies auto_context, then FlowOptions
                 def auto_context_decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
                     wrapped = _apply_auto_context(fn, parent=context_parent)
@@ -589,6 +587,13 @@ def load_prices(
 
         return flow_model(*args, **kwargs)
 
+    @staticmethod
+    def transform(*args, **kwargs):
+        """Decorator that turns a top-level function into a serializable with_inputs() transform factory."""
+        from .flow_model import flow_transform
+
+        return flow_transform(*args, **kwargs)
+
 
 # *****************************************************************************
 # Define "Evaluators" and associated types
@@ -965,31 +970,27 @@ def __call__(self, *, x: int, y: str = "default") -> GenericResult:
         model = MyCallable()
         model(x=42, y="hello")  # Works with kwargs
     """
-    sig = signature(func)
+    from .flow_model import _callable_qualname, _resolved_flow_signature
+
+    sig = _resolved_flow_signature(
+        func,
+        skip_self=True,
+        require_return_annotation=True,
+        annotation_error_suffix=" when auto_context=True",
+        return_error_suffix=" when auto_context=True",
+        function_name=_callable_qualname(func),
+    )
     base_class = parent or ContextBase
 
-    if sig.return_annotation is inspect.Signature.empty:
-        raise TypeError(f"Function {_callable_qualname(func)} must have a return type annotation when auto_context=True")
-
     # Validate parent fields are in function signature
     if parent is not None:
         parent_fields = set(parent.model_fields.keys()) - set(ContextBase.model_fields.keys())
-        sig_params = set(sig.parameters.keys()) - {"self"}
+        sig_params = set(sig.parameters)
         missing = parent_fields - sig_params
         if missing:
             raise TypeError(f"Parent context fields {missing} must be included in function signature")
 
-    # Build fields from parameters (skip 'self'), pydantic validates types
-    fields = {}
-    for name, param in sig.parameters.items():
-        if name == "self":
-            continue
-        if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
-            raise TypeError(f"Function {_callable_qualname(func)} does not support {param.kind.description} when auto_context=True")
-        if param.annotation is inspect.Parameter.empty:
-            raise TypeError(f"Parameter '{name}' must have a type annotation when auto_context=True")
-        default = ... if param.default is inspect.Parameter.empty else param.default
-        fields[name] = (param.annotation, default)
+    fields = {name: (param.annotation, ... if param.default is inspect.Parameter.empty else param.default) for name, param in sig.parameters.items()}
 
     # Create auto context class
     auto_context_class = create_ccflow_model(f"{_callable_qualname(func)}_AutoContext", __base__=base_class, **fields)