Don't create dynamic module, add more tests for locally defined models

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

Author: NeejWeej

ccflow/base.py

@@ -30,7 +30,7 @@
 from typing_extensions import Self
 
 from .exttypes.pyobjectpath import PyObjectPath
-from .local_persistence import _register_local_subclass
+from .local_persistence import _register, _register_local_subclass_if_needed, _sync_to_module
 
 log = logging.getLogger(__name__)
 
@@ -157,14 +157,29 @@ 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)
+        # Register local-scope classes so they're importable via PyObjectPath.
+        # At definition time, we can only detect <locals> in qualname - full check deferred.
+        if "__ccflow_import_path__" in cls.__dict__:
+            _sync_to_module(cls)  # Cross-process unpickle case
+        elif "<locals>" in cls.__qualname__:
+            _register(cls)
+
+    def __reduce_ex__(self, protocol):
+        # Ensure the class is registered before pickling. This is necessary for classes
+        # created dynamically (e.g., via pydantic's create_model) that don't have '<locals>'
+        # in their __qualname__. Without this, such classes would get different UUIDs in
+        # different processes if pickled before type_ is accessed, breaking cross-process
+        # consistency. By registering here, we guarantee __ccflow_import_path__ is set
+        # before cloudpickle serializes the class definition.
+        #
+        # We use __reduce_ex__ instead of __reduce__ because it's called first, allowing us
+        # to perform the registration side effect. We then return NotImplemented to let
+        # pydantic's default pickling mechanism take over, preserving the pickle format.
+        _register_local_subclass_if_needed(type(self))
+        return super().__reduce_ex__(protocol)
 
     @computed_field(
         alias="_target_",
@@ -830,8 +845,6 @@ 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,8 +74,6 @@ 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/exttypes/pyobjectpath.py

@@ -1,44 +1,86 @@
 """This module contains extension types for pydantic."""
 
+import importlib
 from functools import cached_property, lru_cache
 from types import FunctionType, MethodType, ModuleType
 from typing import Any, Type, get_origin
 
-from pydantic import ImportString, TypeAdapter
+from pydantic import TypeAdapter
 from pydantic_core import core_schema
 from typing_extensions import Self
 
-from ccflow.local_persistence import _ensure_registered_at_import_path
-
-_import_string_adapter = TypeAdapter(ImportString)
+from ccflow.local_persistence import _register_local_subclass_if_needed
 
 
 @lru_cache(maxsize=None)
-def import_string(input_string: str):
-    return _import_string_adapter.validate_python(input_string)
+def import_string(dotted_path: str) -> Any:
+    """Import an object from a dotted path string.
+
+    Handles nested class paths like 'module.OuterClass.InnerClass' by progressively
+    trying shorter module paths and using getattr for the remaining parts.
+
+    This is more flexible than pydantic's ImportString which can fail on nested classes.
+    """
+    if not dotted_path:
+        raise ImportError("Empty path")
+
+    parts = dotted_path.split(".")
+
+    # Try progressively shorter module paths
+    # e.g., for 'a.b.C.D', try 'a.b.C', then 'a.b', then 'a'
+    for i in range(len(parts), 0, -1):
+        module_path = ".".join(parts[:i])
+        try:
+            obj = importlib.import_module(module_path)
+            # Successfully imported module, now getattr for remaining parts
+            for attr_name in parts[i:]:
+                obj = getattr(obj, attr_name)
+            return obj
+        except ImportError:
+            continue
+        except AttributeError:
+            # Module imported but attribute not found - keep trying shorter paths
+            continue
+
+    raise ImportError(f"No module named '{dotted_path}'")
+
+
+def _build_standard_import_path(obj: Any) -> str:
+    """Build 'module.qualname' path from an object with __module__ and __qualname__."""
+    qualname = obj.__qualname__
+    # Strip generic type parameters (e.g., "MyClass[int]" -> "MyClass")
+    # This happens with Generic types in pydantic. Type info is lost but
+    # at least the base class remains importable.
+    # TODO: Find a way of capturing the underlying type info
+    if "[" in qualname:
+        qualname = qualname.split("[", 1)[0]
+    return f"{obj.__module__}.{qualname}" if obj.__module__ else qualname
 
 
 class PyObjectPath(str):
-    """Similar to pydantic's ImportString (formerly PyObject in v1), this class represents the path to the object as a string.
+    """A string representing an importable Python object path (e.g., "module.ClassName").
 
-    In pydantic v1, PyObject could not be serialized to json, whereas in v2, ImportString can.
-    However, the round trip is not always consistent, i.e.
+    Similar to pydantic's ImportString, but with consistent serialization behavior:
+    - ImportString deserializes to the actual object
+    - PyObjectPath deserializes back to the string path
+
+    Example:
         >>> ta = TypeAdapter(ImportString)
         >>> ta.validate_json(ta.dump_json("math.pi"))
         3.141592653589793
         >>> ta = TypeAdapter(PyObjectPath)
         >>> ta.validate_json(ta.dump_json("math.pi"))
         'math.pi'
 
-    Other differences are that ImportString can contain other arbitrary python values, whereas PyObjectPath is always a string
+    PyObjectPath also only accepts importable objects, not arbitrary values:
         >>> TypeAdapter(ImportString).validate_python(0)
         0
         >>> TypeAdapter(PyObjectPath).validate_python(0)
         raises
     """
 
     # TODO: It would be nice to make this also derive from Generic[T],
-    #  where T could then by used for type checking in validate.
+    #  where T could then be used for type checking in validate.
     #  However, this doesn't work: https://github.com/python/typing/issues/629
 
     @cached_property
@@ -52,41 +94,41 @@ def __get_pydantic_core_schema__(cls, source_type, handler):
 
     @classmethod
     def _validate(cls, value: Any):
+        """Convert value (string path or object) to PyObjectPath, verifying it's importable."""
         if isinstance(value, str):
-            value = cls(value)
-        else:  # Try to construct a string from the object that can then be used to import the object
+            path = cls(value)
+        else:
+            # Unwrap generic types (e.g., List[int] -> list)
             origin = get_origin(value)
             if origin:
                 value = origin
+            path = cls._path_from_object(value)
 
-            # Check for ccflow's import path override first (used for local-scope classes)
-            # This allows classes with '<locals>' in __qualname__ to remain importable
-            # while preserving cloudpickle's ability to serialize the class definition
-            if hasattr(value, "__ccflow_import_path__"):
-                _ensure_registered_at_import_path(value)
-                value = cls(value.__ccflow_import_path__)
-            elif hasattr(value, "__module__") and hasattr(value, "__qualname__"):
-                if value.__module__ == "__builtin__":
-                    module = "builtins"
-                else:
-                    module = value.__module__
-                qualname = value.__qualname__
-                if "[" in qualname:
-                    # This happens with Generic types in pydantic. We strip out the info for now.
-                    # TODO: Find a way of capturing the underlying type info
-                    qualname = qualname.split("[", 1)[0]
-                if not module:
-                    value = cls(qualname)
-                else:
-                    value = cls(module + "." + qualname)
-            else:
-                raise ValueError(f"ensure this value contains valid import path or importable object: unable to import path for {value}")
+        # Verify the path is actually importable
         try:
-            value.object
+            path.object
         except ImportError as e:
             raise ValueError(f"ensure this value contains valid import path or importable object: {str(e)}")
 
-        return value
+        return path
+
+    @classmethod
+    def _path_from_object(cls, value: Any) -> "PyObjectPath":
+        """Build import path from an object. Triggers ccflow registration for local classes."""
+        if isinstance(value, type):
+            # For ccflow BaseModel subclasses that aren't normally importable (defined in
+            # functions or via create_model), this registers them on ccflow.base
+            _register_local_subclass_if_needed(value)
+
+            # Use __ccflow_import_path__ if set (check __dict__ to avoid inheriting from parents)
+            if "__ccflow_import_path__" in value.__dict__:
+                return cls(value.__ccflow_import_path__)
+            return cls(_build_standard_import_path(value))
+
+        if hasattr(value, "__module__") and hasattr(value, "__qualname__"):
+            return cls(_build_standard_import_path(value))
+
+        raise ValueError(f"ensure this value contains valid import path or importable object: unable to import path for {value}")
 
     @classmethod
     @lru_cache(maxsize=None)
@@ -95,10 +137,12 @@ def _validate_cached(cls, value: str) -> Self:
 
     @classmethod
     def validate(cls, value) -> Self:
-        """Try to convert/validate an arbitrary value to a PyObjectPath."""
-        if isinstance(
-            value, (str, type, FunctionType, ModuleType, MethodType)
-        ):  # If the value is trivial, we cache it here to avoid the overhead of validation
+        """Try to convert/validate an arbitrary value to a PyObjectPath.
+
+        Uses caching for common value types to improve performance.
+        """
+        # Cache validation for common immutable types to avoid repeated work
+        if isinstance(value, (str, type, FunctionType, ModuleType, MethodType)):
             return cls._validate_cached(value)
         return _TYPE_ADAPTER.validate_python(value)