pydantic + sqlmodel namespace hooks for lazy forward-ref resolution

fquery/pydantic.py

@@ -1,17 +1,93 @@
 import dataclasses
+import re
 from dataclasses import dataclass, fields
-from typing import Type
+from typing import Callable, Mapping, Type
 
 from pydantic import BaseModel, ConfigDict, Field
 
+# Matches "name 'Foo' is not defined" raised for unresolvable forward refs.
+_MISSING = re.compile(r"name '(\w+)' is not defined")
 
-def pydantic(cls):
-    return model(dataclass(kw_only=True)(cls))
+# A namespace is either a static mapping or a callable taking the decorated
+# class and returning a mapping. Callables allow consumers (e.g. generated
+# packages with hundreds of lazily-imported models) to supply forward-ref
+# targets on demand instead of importing everything up front.
+Namespace = Mapping | Callable[[type], Mapping]
+
+
+def pydantic(cls=None, *, namespace: Namespace | None = None):
+    """Decorate a class as an fquery pydantic dataclass.
+
+    Usable bare (``@pydantic``) or parametrized
+    (``@pydantic(namespace=...)``). ``namespace`` supplies extra names for
+    resolving forward references when ``validator()`` rebuilds the model;
+    it is inherited by subclasses unless they pass their own.
+    """
+    if cls is None:
+        return lambda c: model(dataclass(kw_only=True)(c), namespace=namespace)
+    return model(dataclass(kw_only=True)(cls), namespace=namespace)
+
+
+def _seed_namespace(cls: type) -> dict:
+    """Names always available: the MRO classes plus their module globals."""
+    import sys
+
+    ns = {c.__name__: c for c in cls.__mro__}
+    for c in cls.__mro__:
+        mod = sys.modules.get(getattr(c, "__module__", ""))
+        if mod is not None:
+            for k, v in vars(mod).items():
+                if isinstance(v, type) and k not in ns:
+                    ns[k] = v
+    return ns
+
+
+def _missing_name(ex: NameError) -> str | None:
+    name = getattr(ex, "name", None)
+    if name:
+        return name
+    m = _MISSING.search(str(ex))
+    return m.group(1) if m else None
+
+
+def _call_provider(provider, cls):
+    try:
+        return provider(cls)
+    except TypeError:
+        return provider()
+
+
+def _rebuild(cls: type, m: type[BaseModel], tries: int = 5) -> None:
+    provider = getattr(cls, "__pydantic_namespace__", None)
+    ns = _seed_namespace(cls)
+    if isinstance(provider, Mapping):
+        ns.update(provider)
+        provider = None  # static: single rebuild attempt below suffices
+    for _ in range(tries):
+        if provider is not None:
+            extra = _call_provider(provider, cls)
+            if extra:
+                ns.update(extra)
+        before = len(ns)
+        try:
+            m.model_rebuild(_types_namespace=dict(ns))
+            return
+        except NameError as ex:
+            if provider is None:
+                raise
+            name = _missing_name(ex)
+            # No progress possible: unknown name, or provider is exhausted.
+            if name is None or name in ns or len(ns) == before:
+                raise
+    raise RuntimeError(f"could not resolve forward refs for {cls.__name__}")
 
 
 def validator(self) -> BaseModel:
-    attrs = {name: getattr(self, name) for name in self.__pydantic__.model_fields}
-    return self.__pydantic__(**attrs)
+    m = self.__pydantic__
+    if not getattr(m, "__pydantic_complete__", True):
+        _rebuild(type(self), m)
+    attrs = {name: getattr(self, name) for name in m.model_fields}
+    return m(**attrs)
 
 
 def get_field_def(cls, field):
@@ -24,7 +100,7 @@ def get_field_def(cls, field):
     return Field(**kwargs)
 
 
-def model(cls: Type) -> Type:
+def model(cls: Type, namespace: Namespace | None = None) -> Type:
     """
     Decorator to convert a dataclass to a Pydantic model.
     """
@@ -40,6 +116,8 @@ def model(cls: Type) -> Type:
         },
     )
     cls.__pydantic__ = pydantic_cls
+    if namespace is not None:
+        cls.__pydantic_namespace__ = namespace
     cls.model_config = ConfigDict(extra="ignore")
     cls.validator = validator
 

fquery/sqlmodel.py

@@ -2,10 +2,12 @@
 from dataclasses import _FIELD, dataclass, field, fields, is_dataclass
 from datetime import date, datetime, time
 from typing import (
+    Callable,
     ClassVar,
     Dict,
     ForwardRef,
     List,
+    Mapping,
     Optional,
     Union,
     get_args,
@@ -83,16 +85,55 @@ def many_to_one(key_column=None, back_populates=None):
     return ret
 
 
-def sqlmodel(cls):
-    return model()(dataclass(kw_only=True)(cls))
+# A namespace is either a static mapping or a callable taking the decorated
+# class and returning a mapping. Callables allow consumers (e.g. generated
+# packages with hundreds of lazily-imported models) to supply forward-ref
+# targets on demand instead of importing everything up front. Mirrors
+# fquery.pydantic's namespace hook; table creation itself stays eager.
+Namespace = Mapping | Callable[[type], Mapping]
+
+
+def sqlmodel(
+    cls=None,
+    *,
+    table: bool = True,
+    table_name: str = None,
+    global_id: bool = False,
+    namespace: Namespace | None = None,
+):
+    """Decorate a class as an fquery sqlmodel dataclass.
+
+    Usable bare (``@sqlmodel``) or parametrized (``@sqlmodel(table=False,
+    namespace=...)``). ``namespace`` supplies extra names for resolving
+    forward references; it is inherited by subclasses unless they pass
+    their own.
+    """
+    decorator = model(
+        table=table,
+        table_name=table_name,
+        global_id=global_id,
+        namespace=namespace,
+    )
+    if cls is None:
+        return lambda c: decorator(dataclass(kw_only=True)(c))
+    return decorator(dataclass(kw_only=True)(cls))
 
 
-def model(table: bool = True, table_name: str = None, global_id: bool = False):
+def model(
+    table: bool = True,
+    table_name: str = None,
+    global_id: bool = False,
+    namespace: Namespace | None = None,
+):
     """
     A decorator that generates a SQLModel from a dataclass.
 
     Args:
         table_name (str): The name of the database table. Defaults to the name of the dataclass.
+        namespace: Extra names for resolving forward references (static
+            mapping or callable taking the decorated class). Useful when
+            related classes are lazily imported and absent from module
+            globals at decoration time.
 
     Returns:
         A decorator that generates a SQLModel from a dataclass.
@@ -160,6 +201,56 @@ def get_field_def(cls, field) -> Union[Field, Relationship]:
             return Field(default=None, foreign_key=sql_meta["foreign_key"])
         raise "Unsupported case"
 
+    def resolve_namespace(cls):
+        provider = getattr(cls, "__sqlmodel_namespace__", namespace)
+        if callable(provider):
+            try:
+                return provider(cls)
+            except TypeError:
+                return provider()
+        return provider or {}
+
+    def type_hints(cls):
+        import sys as _sys
+
+        globalns = {}
+        for klass in getattr(cls, "__mro__", (cls,)):
+            mod = _sys.modules.get(getattr(klass, "__module__", ""))
+            if mod is not None:
+                for k, v in vars(mod).items():
+                    globalns.setdefault(k, v)
+        extra = resolve_namespace(cls)
+        if extra:
+            globalns.update(extra)
+        return get_type_hints(cls, globalns=globalns)
+
+    def rebuild_model(sqlmodel_cls, cls):
+        provider = getattr(cls, "__sqlmodel_namespace__", namespace)
+        if provider is None:
+            return
+        if getattr(sqlmodel_cls, "__pydantic_complete__", True):
+            return
+        import sys as _sys
+
+        ns = {}
+        for klass in getattr(cls, "__mro__", (cls,)):
+            mod = _sys.modules.get(getattr(klass, "__module__", ""))
+            if mod is not None:
+                for k, v in vars(mod).items():
+                    ns.setdefault(k, v)
+        extra = resolve_namespace(cls)
+        if extra:
+            ns.update(extra)
+        rebuild = getattr(sqlmodel_cls, "sqlmodel_rebuild", None) or getattr(
+            sqlmodel_cls, "model_rebuild", None
+        )
+        if rebuild is None:
+            return
+        try:
+            rebuild(_types_namespace=dict(ns))
+        except TypeError:
+            rebuild()
+
     def get_field_type(field, cls):
         sql_meta = field.metadata.get("SQL", {})
         has_foreign_key = bool(sql_meta.get("foreign_key", None))
@@ -174,12 +265,22 @@ def get_field_type(field, cls):
             other_class = type_class.__args__[0]
             if has_many_to_one_relationship:
                 try:
-                    type_class = get_type_hints(cls)[field.name]
+                    type_class = type_hints(cls)[field.name]
                 except NameError:
                     # TODO: log exception?
                     pass
                 else:
-                    return Optional[other_class.__sqlmodel__]
+                    # Use the resolved target, not the raw annotation:
+                    # the latter may still be a ForwardRef when the field
+                    # was declared as a string (e.g. Optional["Widget"]).
+                    args = get_args(type_class)
+                    target = args[0] if args else other_class
+                    if isinstance(target, ForwardRef):
+                        return field.type
+                    target_sqlmodel = getattr(target, "__sqlmodel__", None)
+                    if target_sqlmodel is None:
+                        return field.type
+                    return Optional[target_sqlmodel]
         return field.type
 
     def patch_back_populates_types(field, back_populates, cls, sqlmodel_cls):
@@ -190,7 +291,7 @@ def patch_back_populates_types(field, back_populates, cls, sqlmodel_cls):
             if has_many_to_one_relationship:
                 type_class = field.type
                 try:
-                    type_class = get_type_hints(cls)[field.name]
+                    type_class = type_hints(cls)[field.name]
                 except NameError:
                     # TODO: log exception?
                     pass
@@ -297,6 +398,11 @@ def decorator(cls):
             table=table,
         )
         cls.__sqlmodel__ = sqlmodel_cls
+        if namespace is not None:
+            cls.__sqlmodel_namespace__ = namespace
+        # Resolve remaining forward refs (e.g. targets supplied via the
+        # namespace hook that are absent from module globals).
+        rebuild_model(sqlmodel_cls, cls)
         # Update type annotations in any class with a relationship with this class to point
         # to the SQLModel, not the dataclass
         for cfield in fields(cls):

pyproject.toml

@@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "fquery"
-version = "0.4"
+version = "0.5"
 description = "A graph query engine"
 readme = "README.md"
 requires-python = ">=3.8"

setup.py

@@ -11,7 +11,7 @@
 
 setuptools.setup(
     name="fquery",
-    version="0.4",
+    version="0.5",
     description="A graph query engine",
     url="https://github.com/adsharma/fquery",
     long_description=long_description,

tests/__init__.py

@@ -1,10 +1,7 @@
 import pathlib
-import sys
 import unittest
 
 PROJECT_DIR = pathlib.Path(__file__).parents[1]
-print(PROJECT_DIR)
-sys.path.append(str(PROJECT_DIR / "fquery"))
 
 
 def test_suite():

tests/test_pydantic.py

@@ -1,4 +1,5 @@
 from dataclasses import is_dataclass
+from typing import Optional
 
 import pytest
 from pydantic import BaseModel, ValidationError
@@ -32,3 +33,36 @@ def test_pydantic_fail():
     u1 = User(name="John Doe", age=42.3)
     with pytest.raises(ValidationError):
         _ = u1.validator()
+
+
+def test_pydantic_namespace_hook_static():
+    # Forward ref resolved from a static mapping, no import of the
+    # target module needed at class-definition time.
+    @pydantic(namespace={"Address": Address})  # noqa: F821
+    class Customer:
+        name: str
+        address: Optional["Address"] = None  # noqa: F821
+
+    c = Customer(name="x")
+    assert isinstance(c.validator(), BaseModel)
+
+
+def test_pydantic_namespace_hook_callable():
+    # Callable receives the decorated class; inherited when not overridden.
+    seen = []
+
+    def provider(cls):
+        seen.append(cls.__name__)
+        return {"Address": Address}
+
+    @pydantic(namespace=provider)
+    class Order:
+        ref: Optional["Address"] = None  # noqa: F821
+
+    assert isinstance(Order().validator(), BaseModel)
+    assert seen == ["Order"]
+
+
+@pydantic
+class Address:
+    street: str = ""