Fix union and coalesce expressions not decoding to the correct type.

Author: dnwpark

gel/_internal/_codegen/_models/_pydantic.py

@@ -5795,6 +5795,45 @@ def resolve(
                             f"# type: ignore [assignment, misc, unused-ignore]"
                         )
 
+                    if function.schemapath in {
+                        SchemaPath('std', 'UNION'),
+                        SchemaPath('std', 'IF'),
+                        SchemaPath('std', '??'),
+                    }:
+                        # Special case for the UNION, IF and ?? operators
+                        # Produce a union type instead of just taking the first
+                        # valid type.
+                        #
+                        # See gel: compile_operator
+                        create_union = self.import_name(
+                            BASE_IMPL, "create_union"
+                        )
+
+                        tvars: list[str] = []
+                        for param, path in sources:
+                            if (
+                                param.name in required_generic_params
+                                or param.name in optional_generic_params
+                            ):
+                                pn = param_vars[param.name]
+                                tvar = f"__t_{pn}__"
+
+                                resolve(pn, path, tvar)
+                                tvars.append(tvar)
+
+                        self.write(
+                            f"{gtvar} = {tvars[0]}  "
+                            f"# type: ignore [assignment, misc, unused-ignore]"
+                        )
+                        for tvar in tvars[1:]:
+                            self.write(
+                                f"{gtvar} = {create_union}({gtvar}, {tvar})  "
+                                f"# type: ignore ["
+                                f"assignment, misc, unused-ignore]"
+                            )
+
+                        continue
+
                     # Try to infer generic type from required params first
                     for param, path in sources:
                         if param.name in required_generic_params:

gel/_internal/_qbmodel/_abstract/__init__.py

@@ -65,6 +65,8 @@
 from ._methods import (
     BaseGelModel,
     BaseGelModelIntersection,
+    BaseGelModelUnion,
+    create_union,
 )
 
 
@@ -132,6 +134,7 @@
     "ArrayMeta",
     "BaseGelModel",
     "BaseGelModelIntersection",
+    "BaseGelModelUnion",
     "ComputedLinkSet",
     "ComputedLinkWithPropsSet",
     "ComputedMultiLinkDescriptor",
@@ -174,6 +177,7 @@
     "TupleMeta",
     "UUIDImpl",
     "copy_or_ref_lprops",
+    "create_union",
     "empty_set_if_none",
     "field_descriptor",
     "get_base_scalars_backed_by_py_type",

gel/_internal/_qbmodel/_abstract/_methods.py

@@ -19,6 +19,7 @@
 from gel._internal import _qb
 from gel._internal._schemapath import (
     TypeNameIntersection,
+    TypeNameUnion,
 )
 from gel._internal import _type_expression
 from gel._internal._xmethod import classonlymethod
@@ -256,6 +257,17 @@ class BaseGelModelIntersection(
     rhs: ClassVar[type[AbstractGelModel]]
 
 
+class BaseGelModelUnion(
+    BaseGelModel,
+    _type_expression.Union,
+    Generic[_T_Lhs, _T_Rhs],
+):
+    __gel_type_class__: ClassVar[type]
+
+    lhs: ClassVar[type[AbstractGelModel]]
+    rhs: ClassVar[type[AbstractGelModel]]
+
+
 T = TypeVar('T')
 U = TypeVar('U')
 
@@ -429,3 +441,93 @@ def process_path_alias(
     _type_intersection_cache[lhs][rhs] = result
 
     return result
+
+
+_type_union_cache: weakref.WeakKeyDictionary[
+    type[AbstractGelModel],
+    weakref.WeakKeyDictionary[
+        type[AbstractGelModel],
+        type[
+            BaseGelModelUnion[type[AbstractGelModel], type[AbstractGelModel]]
+        ],
+    ],
+] = weakref.WeakKeyDictionary()
+
+
+def create_union(
+    lhs: _T_Lhs,
+    rhs: _T_Rhs,
+) -> type[BaseGelModelUnion[_T_Lhs, _T_Rhs]]:
+    """Create a runtime union type which acts like a GelModel."""
+
+    if (lhs_entry := _type_union_cache.get(lhs)) and (
+        rhs_entry := lhs_entry.get(rhs)
+    ):
+        return rhs_entry  # type: ignore[return-value]
+
+    # Combine pointer reflections from args
+    ptr_reflections: dict[str, _qb.GelPointerReflection] = {
+        p_name: p_refl
+        for p_name, p_refl in lhs.__gel_reflection__.pointers.items()
+        if p_name in rhs.__gel_reflection__.pointers
+    }
+
+    # Create type reflection for union type
+    class __gel_reflection__(_qb.GelObjectTypeExprMetadata.__gel_reflection__):  # noqa: N801
+        expr_object_types: set[type[AbstractGelModel]] = getattr(
+            lhs.__gel_reflection__, 'expr_object_types', {lhs}
+        ) | getattr(rhs.__gel_reflection__, 'expr_object_types', {rhs})
+
+        type_name = TypeNameUnion(
+            args=(
+                lhs.__gel_reflection__.type_name,
+                rhs.__gel_reflection__.type_name,
+            )
+        )
+
+        pointers = ptr_reflections
+
+        @classmethod
+        def object(
+            cls,
+        ) -> Any:
+            raise NotImplementedError(
+                "Type expressions schema objects are inaccessible"
+            )
+
+    result = type(
+        f"({lhs.__name__} | {rhs.__name__})",
+        (BaseGelModelUnion,),
+        {
+            'lhs': lhs,
+            'rhs': rhs,
+            '__gel_reflection__': __gel_reflection__,
+        },
+    )
+
+    # Generate path aliases for pointers.
+    #
+    # These are used to generate the appropriate path prefix when getting
+    # pointers in shapes.
+    path_aliases: dict[str, _qb.PathAlias] = {
+        p_name: l_path_alias
+        for p_name, p_refl in lhs.__gel_reflection__.pointers.items()
+        if (
+            hasattr(lhs, p_name)
+            and (l_path_alias := getattr(lhs, p_name, None)) is not None
+            and isinstance(l_path_alias, _qb.PathAlias)
+        )
+        if (
+            hasattr(rhs, p_name)
+            and (r_path_alias := getattr(rhs, p_name, None)) is not None
+            and isinstance(r_path_alias, _qb.PathAlias)
+        )
+    }
+    for p_name, path_alias in path_aliases.items():
+        setattr(result, p_name, path_alias)
+
+    if lhs not in _type_union_cache:
+        _type_union_cache[lhs] = weakref.WeakKeyDictionary()
+    _type_union_cache[lhs][rhs] = result
+
+    return result

gel/_internal/_typing_dispatch.py

@@ -70,6 +70,8 @@ def _issubclass(lhs: Any, tp: Any, fn: Any) -> bool:
 
     if issubclass(lhs, _type_expression.Intersection):
         return any(_issubclass(c, tp, fn) for c in (lhs.lhs, lhs.rhs))
+    elif issubclass(lhs, _type_expression.Union):
+        return all(_issubclass(c, tp, fn) for c in (lhs.lhs, lhs.rhs))
 
     if _typing_inspect.is_generic_alias(tp):
         origin = typing.get_origin(tp)

gel/models/pydantic.py

@@ -68,6 +68,7 @@
     PyTypeScalarConstraint,
     RangeMeta,
     UUIDImpl,
+    create_union,
     empty_set_if_none,
 )
 
@@ -195,6 +196,7 @@
     "classonlymethod",
     "computed_field",
     "construct_infix_op_chain",
+    "create_union",
     "dispatch_overload",
     "empty_set_if_none",
 )