Add extra stuff, need clean-up

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

Author: Nijat Khanbabayev

ccflow/__init__.py

@@ -12,6 +12,7 @@
 from .context import *
 from .dep import *
 from .enums import Enum
+from .flow_model import FlowAPI, BoundModel, Lazy
 from .global_state import *
 from .local_persistence import *
 from .models import *

ccflow/callable.py

@@ -312,7 +312,10 @@ def _resolve_deps_and_call(model, context, fn):
     # Get Dep-annotated fields for this model class
     dep_fields = _get_dep_fields(model.__class__)
 
-    if not dep_fields:
+    # Check if model has custom deps (from @func.deps decorator)
+    has_custom_deps = getattr(model.__class__, "__has_custom_deps__", False)
+
+    if not dep_fields and not has_custom_deps:
         return fn(model, context)
 
     # Get dependencies from __deps__
@@ -324,27 +327,37 @@ def _resolve_deps_and_call(model, context, fn):
 
     # Resolve dependencies and store in context var
     resolved_values = {}
-    for field_name, dep in dep_fields.items():
-        field_value = getattr(model, field_name, None)
-        if field_value is None:
-            continue
-
-        # Check if field is a CallableModel that needs resolution
-        if not isinstance(field_value, _CallableModel):
-            continue  # Already a resolved value, skip
 
-        # Check if this field is in __deps__ (for custom transforms)
-        if id(field_value) in dep_map:
-            dep_model, contexts = dep_map[id(field_value)]
-            # Call dependency with the (transformed) context
+    # If custom deps, resolve ALL CallableModel fields from dep_map
+    if has_custom_deps:
+        for dep_model, contexts in deps_result:
             resolved = dep_model(contexts[0]) if contexts else dep_model(context)
-        else:
-            # Not in __deps__, use Dep annotation transform directly
-            transformed_ctx = dep.apply(context)
-            resolved = field_value(transformed_ctx)
+            # Unwrap GenericResult if present (consistent with auto-detected deps)
+            if hasattr(resolved, 'value'):
+                resolved = resolved.value
+            resolved_values[id(dep_model)] = resolved
+    else:
+        # Standard path: iterate over Dep-annotated fields
+        for field_name, dep in dep_fields.items():
+            field_value = getattr(model, field_name, None)
+            if field_value is None:
+                continue
+
+            # Check if field is a CallableModel that needs resolution
+            if not isinstance(field_value, _CallableModel):
+                continue  # Already a resolved value, skip
+
+            # Check if this field is in __deps__ (for custom transforms)
+            if id(field_value) in dep_map:
+                dep_model, contexts = dep_map[id(field_value)]
+                # Call dependency with the (transformed) context
+                resolved = dep_model(contexts[0]) if contexts else dep_model(context)
+            else:
+                # Not in __deps__, use Dep annotation transform directly
+                transformed_ctx = dep.apply(context)
+                resolved = field_value(transformed_ctx)
 
-        # Store resolved value keyed by the CallableModel's id
-        resolved_values[id(field_value)] = resolved
+            resolved_values[id(field_value)] = resolved
 
     # Store in context var and call function
     current_store = _resolved_deps.get()

ccflow/context.py

@@ -2,10 +2,10 @@
 
 import warnings
 from datetime import date, datetime
-from typing import Generic, Hashable, Optional, Sequence, Set, TypeVar
+from typing import Any, Generic, Hashable, Optional, Sequence, Set, TypeVar
 
 from deprecated import deprecated
-from pydantic import field_validator, model_validator
+from pydantic import ConfigDict, field_validator, model_validator
 
 from .base import ContextBase
 from .exttypes import Frequency
@@ -15,6 +15,7 @@
 
 
 __all__ = (
+    "FlowContext",
     "NullContext",
     "GenericContext",
     "DateContext",
@@ -93,6 +94,42 @@
 # Starting 0.8.0 Nullcontext is an alias to ContextBase
 NullContext = ContextBase
 
+
+class FlowContext(ContextBase):
+    """Universal context for @Flow.model functions.
+
+    Instead of generating a new ContextBase subclass for each @Flow.model,
+    this single class with extra="allow" serves as the universal carrier.
+    Validation happens via TypedDict + TypeAdapter at compute() time.
+
+    This design avoids:
+    - Proliferation of dynamic _funcname_Context classes
+    - Class registration overhead for serialization
+    - Pickling issues with Ray/distributed computing
+
+    Fields are stored in __pydantic_extra__ and accessed via __getattr__.
+    """
+
+    model_config = ConfigDict(extra="allow", frozen=True)
+
+    def __getattr__(self, name: str) -> Any:
+        """Access fields stored in __pydantic_extra__."""
+        # Use object.__getattribute__ to avoid infinite recursion
+        try:
+            extra = object.__getattribute__(self, "__pydantic_extra__")
+            if extra is not None and name in extra:
+                return extra[name]
+        except AttributeError:
+            pass
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
+
+    def __repr__(self) -> str:
+        """Show all fields including extra fields."""
+        extra = object.__getattribute__(self, "__pydantic_extra__") or {}
+        fields = ", ".join(f"{k}={v!r}" for k, v in extra.items())
+        return f"FlowContext({fields})"
+
+
 C = TypeVar("C", bound=Hashable)