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
@@ -631,7 +636,16 @@ def _context_validator(cls, values: Any, handler: Any, info: Any):
         if isinstance(values, dict):
             model = values.get("model")
             if model and isinstance(model, CallableModel) and not isinstance(values.get("context"), model.context_type):
-                values["context"] = model.context_type.model_validate(values.get("context"))
+                ctx_type = model.context_type
+                ctx_value = values.get("context")
+                # Handle Optional[ContextType] — if context is None, keep it; otherwise validate through the inner type
+                if get_origin(ctx_type) is Union and type(None) in get_args(ctx_type):
+                    if ctx_value is not None:
+                        inner_type = [t for t in get_args(ctx_type) if t is not type(None)][0]
+                        if not isinstance(ctx_value, inner_type):
+                            values["context"] = inner_type.model_validate(ctx_value)
+                else:
+                    values["context"] = ctx_type.model_validate(ctx_value)
 
         # Apply standard pydantic validation
         context = handler(values)
@@ -965,31 +979,41 @@ 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)
+        # Validate parent field type compatibility
+        from .flow_model import _context_type_annotations_compatible
+
+        for fname in parent_fields:
+            parent_annotation = parent.model_fields[fname].annotation
+            func_annotation = sig.parameters[fname].annotation
+            if func_annotation is inspect.Parameter.empty:
+                continue
+            if not _context_type_annotations_compatible(func_annotation, parent_annotation):
+                raise TypeError(
+                    f"auto_context field '{fname}' has annotation {func_annotation!r} which is incompatible "
+                    f"with parent field annotation {parent_annotation!r}"
+                )
+
+    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)

ccflow/context.py

@@ -5,7 +5,7 @@
 from typing import Any, Generic, Hashable, Optional, Sequence, Set, TypeVar
 
 from deprecated import deprecated
-from pydantic import ConfigDict, field_validator, model_validator
+from pydantic import ConfigDict, PrivateAttr, field_validator, model_validator
 
 from .base import ContextBase
 from .exttypes import Frequency
@@ -106,19 +106,30 @@ class FlowContext(ContextBase):
     """
 
     model_config = ConfigDict(extra="allow", frozen=True)
+    _frozen_hash_key: Hashable | None = PrivateAttr(default=None)
+    _hash_value: int | None = PrivateAttr(default=None)
+
+    def _hash_key(self) -> Hashable:
+        if self._frozen_hash_key is None:
+            self._frozen_hash_key = _freeze_for_hash(self.model_dump(mode="python"))
+        return self._frozen_hash_key
 
     def __eq__(self, other: Any) -> bool:
+        if self is other:
+            return True
         if not isinstance(other, FlowContext):
             return False
-        return self.model_dump(mode="python") == other.model_dump(mode="python")
+        return self._hash_key() == other._hash_key()
 
     def __hash__(self) -> int:
-        return hash(_freeze_for_hash(self.model_dump(mode="python")))
+        if self._hash_value is None:
+            self._hash_value = hash(self._hash_key())
+        return self._hash_value
 
 
 def _freeze_for_hash(value: Any) -> Hashable:
     if isinstance(value, Mapping):
-        return tuple(sorted((key, _freeze_for_hash(item)) for key, item in value.items()))
+        return tuple(sorted(((key, _freeze_for_hash(item)) for key, item in value.items()), key=lambda item: repr(item[0])))
     if isinstance(value, (list, tuple)):
         return tuple(_freeze_for_hash(item) for item in value)
     if isinstance(value, (set, frozenset)):
@@ -130,7 +141,7 @@ def _freeze_for_hash(value: Any) -> Hashable:
     except TypeError as exc:
         if hasattr(value, "__dict__"):
             return (type(value), _freeze_for_hash(vars(value)))
-        raise TypeError(f"FlowContext contains an unhashable value of type {type(value).__name__}") from exc
+        raise TypeError(f"FlowContext contains an unhashable value of type {type(value).__name__}: {value!r}") from exc
     return value