expand pydantic utils to auto reload parent modules when new registration happens

Signed-off-by: Mark Kurtz <[email protected]>

Author: markurtz

src/speculators/utils/pydantic_utils.py

@@ -11,7 +11,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any, ClassVar, Generic, TypeVar
+from typing import Any, ClassVar, Generic, TypeVar, get_args, get_origin
 
 from pydantic import BaseModel, GetCoreSchemaHandler
 from pydantic_core import CoreSchema, core_schema
@@ -41,15 +41,93 @@ class ReloadableBaseModel(BaseModel):
     """
 
     @classmethod
-    def reload_schema(cls) -> None:
+    def reload_schema(cls, parents: bool = True) -> None:
         """
         Reload the class schema with updated registry information.
 
         Forces a complete rebuild of the Pydantic model schema to incorporate
         any changes made to associated registries or validation rules.
+
+        :param parents: Whether to also rebuild schemas for any pydantic parent
+            types that reference this model.
         """
         cls.model_rebuild(force=True)
 
+        if parents:
+            cls.reload_parent_schemas()
+
+    @classmethod
+    def reload_parent_schemas(cls):
+        """
+        Recursively reload schemas for all parent Pydantic models.
+
+        Traverses the inheritance hierarchy to find all parent classes that
+        are Pydantic models and triggers schema rebuilding on each to ensure
+        that any changes in child models are reflected in parent schemas.
+        """
+        potential_parents: set[type[BaseModel]] = {BaseModel}
+        stack: list[type[BaseModel]] = [BaseModel]
+
+        while stack:
+            current = stack.pop()
+            for subclass in current.__subclasses__():
+                if (
+                    issubclass(subclass, BaseModel)
+                    and subclass is not cls
+                    and subclass not in potential_parents
+                ):
+                    potential_parents.add(subclass)
+                    stack.append(subclass)
+
+        for check in cls.__mro__:
+            if isinstance(check, type) and issubclass(check, BaseModel):
+                cls._reload_schemas_depending_on(check, potential_parents)
+
+    @classmethod
+    def _reload_schemas_depending_on(cls, target: type[BaseModel], types: set[type]):
+        changed = True
+        while changed:
+            changed = False
+            for candidate in types:
+                if (
+                    isinstance(candidate, type)
+                    and issubclass(candidate, BaseModel)
+                    and any(
+                        cls._uses_type(target, field_info.annotation)
+                        for field_info in candidate.model_fields.values()
+                        if field_info.annotation is not None
+                    )
+                ):
+                    try:
+                        before = candidate.model_json_schema()
+                    except Exception:  # noqa: BLE001
+                        before = None
+                    candidate.model_rebuild(force=True)
+                    if before is not None:
+                        after = candidate.model_json_schema()
+                        changed |= before != after
+
+    @classmethod
+    def _uses_type(cls, target: type, candidate: type) -> bool:
+        if target is candidate:
+            return True
+
+        origin = get_origin(candidate)
+
+        if origin is None:
+            return isinstance(candidate, type) and issubclass(candidate, target)
+
+        if isinstance(origin, type) and (
+            target is origin or issubclass(origin, target)
+        ):
+            return True
+
+        for arg in get_args(candidate) or []:
+            if isinstance(arg, type) and cls._uses_type(target, arg):
+                return True
+
+        return False
+
 
 class PydanticClassRegistryMixin(
     ReloadableBaseModel, RegistryMixin[type[BaseModelT]], ABC, Generic[BaseModelT]