Embed __tname__ in models, support using it in model_validate (#918)

Gel supports inheritance much more than pydantic does, and so when we
have a link to a parent type, we need to be able to determine which of
its potential children type to create.

We do this by adding a `tname__` field to every generated class,
renamed to `__tname__` via the alias system, and update
`__gel_validate__` (which currently just directly calls
`model_validate`) to discriminate on it. The hard part here is
handling the case where there are linkprops, which are represented by
per-link subclasses of ProxyModel.  To deal with this, we dynamically
generate subclasses of the link class.

This mechanism also allows us to return correct types from the client
for links with link properties that point to subtypes, which we
previously weren't.

The other potential option was to use pydantic's discriminated unions,
either by directly generating them in the source or dynamically
mocking them up when creating models.  Getting this to play nicely
with everything we do seemed pretty nasty. (The linkprop situation in
particular scares me more.) The big potential advantage of doing it
that way, instead of having our own custom callbacks, is that it might
allow us to not lose track of the parameters of validation (like
strict mode and whether we are in JSON mode).

But it turns out we were already losing track of that due to
various ways we were interposing ourselves in validation, so
maybe it's fine. We can maybe fix this with contextvars later.
(Though the json situation is rough.)

Fixes #755.

Author: msullivan

gel/_internal/_codegen/_models/_pydantic.py

@@ -4716,6 +4716,13 @@ def write_object_type(
         ):
             self.write(f'"""type {objtype.name}"""')
             self.write()
+            literal = self.import_name("typing", "Literal")
+            field = self.import_name("pydantic", "Field")
+            self._write_model_attribute(
+                "tname__",
+                f'{literal}["{type_name}"] = {field}('
+                f'"{type_name}", alias="__tname__")',
+            )
             pointers = _get_object_type_body(objtype)
             if pointers:
                 localns = frozenset(ptr.name for ptr in pointers)

gel/_internal/_qbmodel/_abstract/__init__.py

@@ -15,6 +15,7 @@
     DefaultValue,
     GelType,
     GelTypeMeta,
+    LITERAL_TAG_FIELDS,
 )
 
 from ._descriptors import (
@@ -110,6 +111,7 @@
 
 __all__ = (
     "DEFAULT_VALUE",
+    "LITERAL_TAG_FIELDS",
     "MODEL_SUBSTRATE_MODULE",
     "AbstractGelLinkModel",
     "AbstractGelModel",

gel/_internal/_qbmodel/_abstract/_base.py

@@ -35,6 +35,9 @@
 T_co = TypeVar("T_co", covariant=True)
 
 
+LITERAL_TAG_FIELDS = ('tname__',)
+
+
 if TYPE_CHECKING:
 
     class GelTypeMeta(abc.ABCMeta):
@@ -162,8 +165,13 @@ def __new__(
             super().__new__(mcls, name, bases, namespace, **kwargs),
         )
         reflection = cls.__gel_reflection__
-        if (tname := getattr(reflection, "name", None)) is not None:
-            mcls.__gel_class_registry__[tname] = cls
+        if (
+            # The class registry only tracks the canonical base instances,
+            # which are the ones that directly declare 'tname__'
+            (tname := getattr(reflection, "name", None)) is not None
+            and 'tname__' in namespace
+        ):
+            mcls.__gel_class_registry__[str(tname)] = cls
         cls.__gel_shape__ = __gel_shape__
         return cls
 

gel/_internal/_qbmodel/_abstract/_descriptors.py

@@ -36,6 +36,7 @@
     AbstractGelLinkModel,
     is_gel_type,
     maybe_collapse_object_type_variant_union,
+    LITERAL_TAG_FIELDS,
 )
 
 
@@ -116,7 +117,9 @@ def _try_resolve_type(self) -> Any:
         if (
             t is not None
             and _typing_inspect.is_generic_alias(t)
-            and issubclass(typing.get_origin(t), PointerDescriptor)
+            and (origin := typing.get_origin(t))
+            and isinstance(origin, type)
+            and issubclass(origin, PointerDescriptor)
         ):
             self.__gel_resolved_descriptor__ = t
             t = typing.get_args(t)[0]
@@ -127,12 +130,15 @@ def _try_resolve_type(self) -> Any:
                 if collapsed is not None:
                     t = collapsed
 
-            if not is_gel_type(t):
+            if (
+                not is_gel_type(t)
+                and self.__gel_name__ not in LITERAL_TAG_FIELDS
+            ):
                 raise AssertionError(
                     f"{self._fqname} type argument is not a GelType: {t}"
                 )
 
-            self.__gel_resolved_type__ = t
+            self.__gel_resolved_type__ = cast('type[GelType]', t)
 
         return t
 
@@ -606,6 +612,9 @@ def get(
 
 class AbstractGelProxyModel(AbstractGelModel, Generic[_MT_co, _LM_co]):
     __linkprops__: GelLinkModelDescriptor[_LM_co]
+    __gel_dynamic_proxy_base__: ClassVar[
+        type[AbstractGelProxyModel[Any, Any]] | None
+    ] = None
 
     if TYPE_CHECKING:
         _p__obj__: _MT_co
@@ -828,7 +837,7 @@ def proxy_link(
 ) -> AbstractGelProxyModel[_MT_co, _LM_co]:
     tp_new = type(new)
 
-    if tp_new is proxy_type:
+    if tp_new is proxy_type or issubclass(tp_new, proxy_type):
         # Fast path for the same proxy type.
 
         new_proxy = cast("AbstractGelProxyModel[_MT_co, _LM_co]", new)

gel/_internal/_qbmodel/_pydantic/_fields.py

@@ -436,7 +436,7 @@ def _validate(
                 )
 
         # defer to Pydantic
-        return mt.model_validate(value)
+        return mt.__gel_validate__(value)  # type: ignore[return-value]
 
 
 class _AnyLinkWithProps(Generic[_PT_co, _BMT_co]):

gel/_internal/_qbmodel/_pydantic/_models.py

@@ -18,6 +18,7 @@
     Self,
 )
 
+import copyreg
 import inspect
 import itertools
 import types
@@ -395,7 +396,11 @@ def _process_pydantic_fields(
 
     for fn, field in fields.items():
         ptr = cls.__gel_reflection__.pointers.get(fn)
-        if (ptr is None or ptr.computed) and fn != "__linkprops__":
+        if (
+            (ptr is None or ptr.computed)
+            and fn != "__linkprops__"
+            and fn not in _abstract.LITERAL_TAG_FIELDS
+        ):
             # Regarding `fn != "__linkprops__"`: see MergedModelMeta --
             # it renames `linkprops____` to `__linkprops__` to circumvent
             # Pydantic's restriction on fields starting with `_`.
@@ -588,6 +593,8 @@ def _process_pydantic_fields(
     validate_assignment=True,
     defer_build=True,
     extra="forbid",
+    serialize_by_alias=True,
+    validate_by_name=False,
 )
 
 
@@ -659,13 +666,19 @@ def __gel_model_construct__(cls, __dict__: dict[str, Any] | None) -> Self:
             ll_setattr(self, "__pydantic_private__", None)
             ll_setattr(self, "__gel_changed_fields__", None)
 
+        ll_setattr(self, "tname__", str(cls.__gel_reflection__.name))
+
         if cls.__gel_has_id_field__:
             mid = self.__dict__.get("id", _unset)
             assert mid is not UNSET_UUID
             ll_setattr(self, "__gel_new__", mid is _unset)
 
         return self
 
+    @property
+    def __tname__(self) -> str:
+        return self.tname__  # type: ignore[no-any-return, attr-defined]
+
     @classmethod
     def model_construct(
         cls,
@@ -890,7 +903,24 @@ def __delattr__(self, name: str) -> None:
 
     @classmethod
     def __gel_validate__(cls, value: Any) -> GelSourceModel:
-        return cls.model_validate(value)
+        # __gel_validate__ is called when validating a gel type
+        # reached by a link, and Gel links may point to subtypes of
+        # the declared type.  We dispatch on that ourselves, using a
+        # __tname__ field in the dict.
+        ccls = cls
+        key = '__tname__'
+        if isinstance(value, dict) and key in value:
+            ncls: Any = cls.get_class_by_name(value[key])
+            if ccls is not ncls:
+                # For proxy models, we need an appropriate proxy.
+                if issubclass(ccls, ProxyModel):
+                    ccls = ccls._get_subtype_proxy(ncls)
+                else:
+                    ccls = ncls
+
+        res = ccls.model_validate(value)
+
+        return res
 
 
 class GelModel(
@@ -950,6 +980,7 @@ def __new__(
 
         validator = cls.__pydantic_validator__
         for arg, value in kwargs.items():
+            arg = _pydantic_utils.GEL_ALIASES.get(arg, arg)  # noqa: PLW2901  # ignoring it in ruff.toml failed?
             # Faster than setattr()
             validator.validate_assignment(self, arg, value)
 
@@ -1216,6 +1247,7 @@ def __build_custom_serializer(
                 "__get_pydantic_core_schema__": classmethod(
                     lambda cls, source_type, handler: handler(source_type)
                 ),
+                "__gel_is_custom_serializer__": True,
             },
         )
 
@@ -1363,6 +1395,18 @@ def model_dump(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
         return pydantic.BaseModel.model_dump(self, *args, **kwargs)
 
 
+def _pickle_dynamic_proxy_model(cls: Any) -> Any:
+    # See discussion in _typing_parametric. We do the same tricks as
+    # PickleableClassParametricType, basically.
+    if "__gel_dynamic_proxy_base__" in cls.__dict__:
+        base = cls.__dict__["__gel_dynamic_proxy_base__"]
+        return base._get_subtype_proxy, (cls.__proxy_of__,)
+    else:
+        # If it is not one of our dynamic things, we return the
+        # classname and pickle goes and does something useful.
+        return cls.__qualname__
+
+
 class ProxyModel(
     GelModel,
     _abstract.AbstractGelProxyModel[_MT_co, GelLinkModel],
@@ -1373,6 +1417,12 @@ class ProxyModel(
     if TYPE_CHECKING:
         __gel_proxy_merged_model_cache__: ClassVar[type[_MergedModelBase]]
 
+    # A cache for dynamically generated proxies where the __proxy_of__
+    # class is a *subtype*. Generated by _get_subtype_proxy.
+    __gel_subtype_proxy_cache__: ClassVar[
+        dict[type[GelModel], type[ProxyModel[Any]]] | None
+    ] = None
+
     # NB: __linkprops__ is not in slots because it is managed by
     #     GelLinkModelDescriptor.
 
@@ -1406,9 +1456,56 @@ def __init__(
         # We want ProxyModel to be a trasparent wrapper, so we
         # forward the constructor arguments to the wrapped object.
         wrapped = self.__proxy_of__(**kwargs)
+        assert not isinstance(wrapped, ProxyModel)
         ll_setattr(self, "_p__obj__", wrapped)
         # __linkprops__ is written into __dict__ by GelLinkModelDescriptor
 
+    def __init_subclass__(cls) -> None:
+        super().__init_subclass__()
+
+        # Register a custom reducer for every *metaclass*
+        # of ProxyModels. We need this to deal with our dynamic classes.
+        # Shouldn't be many.
+        copyreg.pickle(
+            type(cls),
+            _pickle_dynamic_proxy_model,  # pyright: ignore [reportArgumentType]
+        )
+
+    @classmethod
+    def _get_subtype_proxy(cls, ncls: type[GelModel]) -> type[Self]:
+        """Generate a subtype that has ncls as its __proxy_of__"""
+
+        if ncls is cls.__proxy_of__:
+            return cls
+
+        if cls.__gel_subtype_proxy_cache__ is None:
+            cls.__gel_subtype_proxy_cache__ = {}
+
+        if cached := cls.__gel_subtype_proxy_cache__.get(ncls):
+            return cast('type[Self]', cached)
+
+        core_schema = ProxyModel.__dict__['__get_pydantic_core_schema__']
+        new_proxy = type(cls)(  # type: ignore[misc]
+            f'{cls.__name__}[{ncls.__name__}]',
+            (ncls, ProxyModel[ncls], cls),  # type: ignore[valid-type]
+            {
+                k: cls.__dict__[k]
+                for k in (
+                    '__module__',
+                    '__qualname__',
+                )
+                if k in cls.__dict__
+            }
+            | {
+                # This gets overridden in some cases, and we need to restore it
+                '__get_pydantic_core_schema__': core_schema,
+                '__gel_dynamic_proxy_base__': cls,
+            },
+        )
+        new_proxy.model_rebuild()
+        cls.__gel_subtype_proxy_cache__[ncls] = new_proxy
+        return cast('type[Self]', new_proxy)
+
     @classmethod
     def link(cls, obj: _MT_co, /, **link_props: Any) -> Self:  # type: ignore [misc]
         proxy_of = ll_type_getattr(cls, "__proxy_of__")
@@ -1439,9 +1536,12 @@ def link(cls, obj: _MT_co, /, **link_props: Any) -> Self:  # type: ignore [misc]
                     f"are allowed, got {type(obj).__name__}",
                 )
 
-        self = cls.__new__(cls)
-        lprops = cls.__linkprops__(**link_props)
+        ncls = cls._get_subtype_proxy(type(obj))
+
+        self = ncls.__new__(ncls)
+        lprops = ncls.__linkprops__(**link_props)
         ll_setattr(self, "__linkprops__", lprops)
+        assert not isinstance(obj, ProxyModel)
         ll_setattr(self, "_p__obj__", obj)
 
         # Treat newly created link props as if they had all their values
@@ -1521,6 +1621,7 @@ def __pydantic_init_subclass__(cls, **kwargs: Any) -> None:
         generic_meta = cls.__pydantic_generic_metadata__
         if generic_meta["origin"] is ProxyModel and generic_meta["args"]:
             cls.__proxy_of__ = generic_meta["args"][0]
+            assert issubclass(cls.__proxy_of__, _abstract.AbstractGelModel)
 
     @classmethod
     def __make_merged_model(cls) -> type[_MergedModelBase]:
@@ -1543,18 +1644,22 @@ def __make_merged_model(cls) -> type[_MergedModelBase]:
         # _MergedModelBase has a custom metaclass, so we must
         # create a common subclass of that and whatever the
         # metaclass of the real type is.
-        metaclass = types.new_class(
-            f"{cls.__name__}Meta",
-            (_MergedModelMeta, type(cls.__proxy_of__)),
-        )
+        pmcls = type(cls.__proxy_of__)
+        if issubclass(pmcls, _MergedModelMeta):
+            metaclass = pmcls
+        else:
+            metaclass = types.new_class(
+                f"{cls.__name__}Meta",
+                (_MergedModelMeta, pmcls),
+            )
 
         merged = cast(
             "type[_MergedModelBase]",
             pydantic.create_model(
                 cls.__name__,
                 __base__=(
-                    _MergedModelBase,
                     cls.__proxy_of__,
+                    _MergedModelBase,
                 ),  # inherit all wrapped fields
                 __config__=DEFAULT_MODEL_CONFIG,
                 __cls_kwargs__={"metaclass": metaclass},
@@ -1632,11 +1737,14 @@ def __gel_proxy_construct__(
         *,
         linked: bool = False,
     ) -> Self:
-        pnv = cls.__gel_model_construct__(None)
+        ncls = cls._get_subtype_proxy(type(obj))
+
+        pnv = ncls.__gel_model_construct__(None)
+        assert not isinstance(obj, ProxyModel)
         ll_setattr(pnv, "_p__obj__", obj)
 
         if type(lprops) is dict:
-            lp_obj = cls.__linkprops__.__gel_model_construct__(lprops)
+            lp_obj = ncls.__linkprops__.__gel_model_construct__(lprops)
         else:
             lp_obj = lprops  # type: ignore [assignment]
 

gel/_internal/_qbmodel/_pydantic/_utils.py

@@ -29,6 +29,13 @@
     from ._models import GelModel
 
 
+# The real field name is tname__, and we need to manually apply the
+# alias in some places.
+GEL_ALIASES = {
+    "__tname__": "tname__",
+}
+
+
 # TypeAliasType does not support recursive types. mypy errors out with:
 #    error: Cannot resolve name "IncEx" (possible cyclic definition)  [misc]
 #
@@ -201,6 +208,14 @@ def massage_model_dump_kwargs(
     except KeyError:
         exclude = kwargs["exclude"] = set()
 
+    # HACK: Handle exclude for internal gel aliases, since pydantic won't.b
+    for from_, to_ in GEL_ALIASES.items():
+        if from_ in exclude:
+            if isinstance(exclude, dict):
+                exclude[to_] = exclude[from_]
+            elif isinstance(exclude, set):
+                exclude.add(to_)
+
     if model.__gel_has_id_field__ and model.__gel_new__:
         # You're generally not supposed to dump an unsaved model,
         # so this branch might be a bit slow -- we can potentially