Register locally defined CallableModels and Contexts in a module to work with PyObjectPath

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

Author: NeejWeej

ccflow/base.py

@@ -30,6 +30,7 @@
 from typing_extensions import Self
 
 from .exttypes.pyobjectpath import PyObjectPath
+from .local_persistence import register_local_subclass
 
 log = logging.getLogger(__name__)
 
@@ -156,6 +157,15 @@ class BaseModel(PydanticBaseModel, _RegistryMixin, metaclass=_SerializeAsAnyMeta
         - Registration by name, and coercion from string name to allow for object re-use from the configs
     """
 
+    __ccflow_local_registration_kind__: ClassVar[str] = "model"
+
+    @classmethod
+    def __pydantic_init_subclass__(cls, **kwargs):
+        # __pydantic_init_subclass__ is the supported hook point once Pydantic finishes wiring the subclass.
+        super().__pydantic_init_subclass__(**kwargs)
+        kind = getattr(cls, "__ccflow_local_registration_kind__", "model")
+        register_local_subclass(cls, kind=kind)
+
     @computed_field(
         alias="_target_",
         repr=False,
@@ -820,6 +830,8 @@ class ContextBase(ResultBase):
     that is an input into another CallableModel.
     """
 
+    __ccflow_local_registration_kind__: ClassVar[str] = "context"
+
     model_config = ConfigDict(
         frozen=True,
         arbitrary_types_allowed=False,

ccflow/callable.py

@@ -74,6 +74,8 @@ class _CallableModel(BaseModel, abc.ABC):
     The purpose of this class is to provide type definitions of context_type and return_type.
     """
 
+    __ccflow_local_registration_kind__: ClassVar[str] = "callable_model"
+
     model_config = ConfigDict(
         ignored_types=(property,),
     )

ccflow/local_persistence.py

@@ -0,0 +1,72 @@
+"""Helpers for persisting BaseModel-derived classes defined inside local scopes."""
+
+from __future__ import annotations
+
+import re
+import sys
+from collections import defaultdict
+from itertools import count
+from types import ModuleType
+from typing import Any, DefaultDict, Type
+
+__all__ = ("LOCAL_ARTIFACTS_MODULE_NAME", "register_local_subclass")
+
+
+LOCAL_ARTIFACTS_MODULE_NAME = "ccflow._local_artifacts"
+_LOCAL_MODULE_DOC = "Auto-generated BaseModel subclasses created outside importable modules."
+
+_SANITIZE_PATTERN = re.compile(r"[^0-9A-Za-z_]")
+_LOCAL_KIND_COUNTERS: DefaultDict[str, count] = defaultdict(lambda: count())
+
+
+def _ensure_module(name: str, doc: str) -> ModuleType:
+    """Ensure the dynamic module exists so import paths remain stable."""
+    module = sys.modules.get(name)
+    if module is None:
+        module = ModuleType(name, doc)
+        sys.modules[name] = module
+        parent_name, _, attr = name.rpartition(".")
+        if parent_name:
+            parent_module = sys.modules.get(parent_name)
+            if parent_module and not hasattr(parent_module, attr):
+                setattr(parent_module, attr, module)
+    return module
+
+
+_LOCAL_ARTIFACTS_MODULE = _ensure_module(LOCAL_ARTIFACTS_MODULE_NAME, _LOCAL_MODULE_DOC)
+
+
+def _needs_registration(cls: Type[Any]) -> bool:
+    module = getattr(cls, "__module__", "")
+    qualname = getattr(cls, "__qualname__", "")
+    return "<locals>" in qualname or module == "__main__"
+
+
+def _sanitize_identifier(value: str, fallback: str) -> str:
+    sanitized = _SANITIZE_PATTERN.sub("_", value or "")
+    sanitized = sanitized.strip("_") or fallback
+    if sanitized[0].isdigit():
+        sanitized = f"_{sanitized}"
+    return sanitized
+
+
+def _build_unique_name(*, kind_slug: str, name_hint: str) -> str:
+    sanitized_hint = _sanitize_identifier(name_hint, "BaseModel")
+    counter = _LOCAL_KIND_COUNTERS[kind_slug]
+    return f"{kind_slug}__{sanitized_hint}__{next(counter)}"
+
+
+def register_local_subclass(cls: Type[Any], *, kind: str = "model") -> None:
+    """Register BaseModel subclasses created in local scopes."""
+    if getattr(cls, "__module__", "").startswith(LOCAL_ARTIFACTS_MODULE_NAME):
+        return
+    if not _needs_registration(cls):
+        return
+
+    name_hint = f"{getattr(cls, '__module__', '')}.{getattr(cls, '__qualname__', '')}"
+    kind_slug = _sanitize_identifier(kind, "model").lower()
+    unique_name = _build_unique_name(kind_slug=kind_slug, name_hint=name_hint)
+    setattr(_LOCAL_ARTIFACTS_MODULE, unique_name, cls)
+    cls.__module__ = _LOCAL_ARTIFACTS_MODULE.__name__
+    cls.__qualname__ = unique_name
+    setattr(cls, "__ccflow_dynamic_origin__", name_hint)

ccflow/tests/evaluators/test_common.py

@@ -1,11 +1,24 @@
 import logging
 from datetime import date
+from typing import ClassVar
 from unittest import TestCase
 
 import pandas as pd
 import pyarrow as pa
 
-from ccflow import DateContext, DateRangeContext, Evaluator, FlowOptionsOverride, ModelEvaluationContext, NullContext
+from ccflow import (
+    CallableModel,
+    ContextBase,
+    DateContext,
+    DateRangeContext,
+    Evaluator,
+    Flow,
+    FlowOptions,
+    FlowOptionsOverride,
+    GenericResult,
+    ModelEvaluationContext,
+    NullContext,
+)
 from ccflow.evaluators import (
     FallbackEvaluator,
     GraphEvaluator,
@@ -16,6 +29,7 @@
     combine_evaluators,
     get_dependency_graph,
 )
+from ccflow.tests.local_helpers import build_nested_graph_chain
 
 from .util import CircularModel, MyDateCallable, MyDateRangeCallable, MyRaisingCallable, NodeModel, ResultModel
 
@@ -473,3 +487,37 @@ def test_graph_evaluator_circular(self):
         evaluator = GraphEvaluator()
         with FlowOptionsOverride(options={"evaluator": evaluator}):
             self.assertRaises(Exception, root, context)
+
+    def test_graph_evaluator_with_local_models_and_cache(self):
+        ParentCls, ChildCls = build_nested_graph_chain()
+        ChildCls.call_count = 0
+        model = ParentCls(child=ChildCls())
+        evaluator = MultiEvaluator(evaluators=[GraphEvaluator(), MemoryCacheEvaluator()])
+        with FlowOptionsOverride(options=FlowOptions(evaluator=evaluator, cacheable=True)):
+            first = model(NullContext())
+            second = model(NullContext())
+        self.assertEqual(first.value, second.value)
+        self.assertEqual(ChildCls.call_count, 1)
+
+
+class TestMemoryCacheEvaluatorLocal(TestCase):
+    def test_memory_cache_handles_local_context_and_callable(self):
+        class LocalContext(ContextBase):
+            value: int
+
+        class LocalModel(CallableModel):
+            call_count: ClassVar[int] = 0
+
+            @Flow.call
+            def __call__(self, context: LocalContext) -> GenericResult:
+                type(self).call_count += 1
+                return GenericResult(value=context.value * 2)
+
+        evaluator = MemoryCacheEvaluator()
+        LocalModel.call_count = 0
+        model = LocalModel()
+        with FlowOptionsOverride(options=FlowOptions(evaluator=evaluator, cacheable=True)):
+            result1 = model(LocalContext(value=5))
+            result2 = model(LocalContext(value=5))
+        self.assertEqual(result1.value, result2.value)
+        self.assertEqual(LocalModel.call_count, 1)

ccflow/tests/local_helpers.py

@@ -0,0 +1,46 @@
+"""Shared helpers for constructing local-scope contexts/models in tests."""
+
+from typing import ClassVar, Tuple, Type
+
+from ccflow import CallableModel, ContextBase, Flow, GenericResult, GraphDepList, NullContext
+
+
+def build_local_callable(name: str = "LocalCallable") -> Type[CallableModel]:
+    class _LocalCallable(CallableModel):
+        @Flow.call
+        def __call__(self, context: NullContext) -> GenericResult:
+            return GenericResult(value="local")
+
+    _LocalCallable.__name__ = name
+    return _LocalCallable
+
+
+def build_local_context(name: str = "LocalContext") -> Type[ContextBase]:
+    class _LocalContext(ContextBase):
+        value: int
+
+    _LocalContext.__name__ = name
+    return _LocalContext
+
+
+def build_nested_graph_chain() -> Tuple[Type[CallableModel], Type[CallableModel]]:
+    class LocalLeaf(CallableModel):
+        call_count: ClassVar[int] = 0
+
+        @Flow.call
+        def __call__(self, context: NullContext) -> GenericResult:
+            type(self).call_count += 1
+            return GenericResult(value="leaf")
+
+    class LocalParent(CallableModel):
+        child: LocalLeaf
+
+        @Flow.call
+        def __call__(self, context: NullContext) -> GenericResult:
+            return self.child(context)
+
+        @Flow.deps
+        def __deps__(self, context: NullContext) -> GraphDepList:
+            return [(self.child, [context])]
+
+    return LocalParent, LocalLeaf

ccflow/tests/test_base.py

@@ -1,9 +1,10 @@
 from typing import Any, Dict, List
-from unittest import TestCase
+from unittest import TestCase, mock
 
-from pydantic import ConfigDict, ValidationError
+from pydantic import BaseModel as PydanticBaseModel, ConfigDict, ValidationError
 
-from ccflow import BaseModel, PyObjectPath
+from ccflow import BaseModel, CallableModel, ContextBase, Flow, GenericResult, NullContext, PyObjectPath
+from ccflow.local_persistence import LOCAL_ARTIFACTS_MODULE_NAME
 
 
 class ModelA(BaseModel):
@@ -106,6 +107,20 @@ def test_type_after_assignment(self):
         self.assertIsInstance(m.type_, PyObjectPath)
         self.assertEqual(m.type_, path)
 
+    def test_pyobjectpath_requires_ccflow_local_registration(self):
+        class PlainLocalModel(PydanticBaseModel):
+            value: int
+
+        with self.assertRaises(ValueError):
+            PyObjectPath.validate(PlainLocalModel)
+
+        class LocalCcflowModel(BaseModel):
+            value: int
+
+        path = PyObjectPath.validate(LocalCcflowModel)
+        self.assertEqual(path.object, LocalCcflowModel)
+        self.assertTrue(str(path).startswith(f"{LOCAL_ARTIFACTS_MODULE_NAME}."))
+
     def test_validate(self):
         self.assertEqual(ModelA.model_validate({"x": "foo"}), ModelA(x="foo"))
         type_ = "ccflow.tests.test_base.ModelA"
@@ -157,3 +172,52 @@ def test_widget(self):
                 {"expanded": True, "root": "bar"},
             ),
         )
+
+
+class TestLocalRegistrationKind(TestCase):
+    def test_base_model_defaults_to_model_kind(self):
+        with mock.patch("ccflow.base.register_local_subclass") as register:
+
+            class LocalModel(BaseModel):
+                value: int
+
+        register.assert_called_once()
+        args, kwargs = register.call_args
+        self.assertIs(args[0], LocalModel)
+        self.assertEqual(kwargs["kind"], "model")
+
+    def test_context_defaults_to_context_kind(self):
+        with mock.patch("ccflow.base.register_local_subclass") as register:
+
+            class LocalContext(ContextBase):
+                value: int
+
+        register.assert_called_once()
+        args, kwargs = register.call_args
+        self.assertIs(args[0], LocalContext)
+        self.assertEqual(kwargs["kind"], "context")
+
+    def test_callable_defaults_to_callable_kind(self):
+        with mock.patch("ccflow.base.register_local_subclass") as register:
+
+            class LocalCallable(CallableModel):
+                @Flow.call
+                def __call__(self, context: NullContext) -> GenericResult:
+                    return GenericResult(value="ok")
+
+        register.assert_called_once()
+        args, kwargs = register.call_args
+        self.assertIs(args[0], LocalCallable)
+        self.assertEqual(kwargs["kind"], "callable_model")
+
+    def test_explicit_override_respected(self):
+        with mock.patch("ccflow.base.register_local_subclass") as register:
+
+            class CustomKind(BaseModel):
+                __ccflow_local_registration_kind__ = "custom"
+                value: int
+
+        register.assert_called_once()
+        args, kwargs = register.call_args
+        self.assertIs(args[0], CustomKind)
+        self.assertEqual(kwargs["kind"], "custom")