Make the return type of type intersection the same as the argument type. (#954)

Author: dnwpark

gel/_internal/_qbmodel/_abstract/_methods.py

@@ -42,8 +42,7 @@
     from collections.abc import Callable
 
 
-_T_SelfModel = TypeVar("_T_SelfModel", bound="type[BaseGelModel]")
-_T_OtherModel = TypeVar("_T_OtherModel", bound="type[BaseGelModel]")
+_T_OtherModel = TypeVar("_T_OtherModel", bound="BaseGelModel")
 
 
 class BaseGelModel(AbstractGelModel):
@@ -84,10 +83,13 @@ def limit(cls, /, expr: Any) -> type[Self]: ...
         @classmethod
         def offset(cls, /, expr: Any) -> type[Self]: ...
 
+        # We pretend that the return type is _T_OtherModel so that the type
+        # checker is aware of _T_OtherModel's pointers. We don't get Self's
+        # pointers, but that's ok most of the time.
         @classmethod
         def is_(
-            cls: _T_SelfModel, /, other_model: _T_OtherModel
-        ) -> type[BaseGelModelIntersection[_T_SelfModel, _T_OtherModel]]: ...
+            cls: type[Self], /, other_model: type[_T_OtherModel]
+        ) -> type[_T_OtherModel]: ...
 
         @classmethod
         def __gel_assert_single__(
@@ -207,11 +209,11 @@ def offset(
         @_qb.exprmethod
         @classmethod
         def is_(
-            cls: _T_SelfModel,
+            cls: type[Self],
             /,
-            value: _T_OtherModel,
+            value: type[_T_OtherModel],
             __operand__: _qb.ExprAlias | None = None,
-        ) -> type[BaseGelModelIntersection[_T_SelfModel, _T_OtherModel]]:
+        ) -> type[BaseGelModelIntersection[type[Self], type[_T_OtherModel]]]:
             return _qb.AnnotatedExpr(  # type: ignore [return-value]
                 create_intersection(cls, value),
                 add_object_type_filter(cls, value, __operand__=__operand__),

gel/_internal/_qbmodel/_pydantic/_models.py

@@ -923,10 +923,6 @@ def __gel_validate__(cls, value: Any) -> GelSourceModel:
         return res
 
 
-_T_SelfModel = TypeVar("_T_SelfModel", bound="type[_abstract.BaseGelModel]")
-_T_OtherModel = TypeVar("_T_OtherModel", bound="type[_abstract.BaseGelModel]")
-
-
 class GelModel(
     GelSourceModel,
     _abstract.BaseGelModel,
@@ -1326,20 +1322,6 @@ def model_copy(
             ll_setattr(copied, "__gel_new__", ll_getattr(self, "__gel_new__"))
         return copied
 
-    if TYPE_CHECKING:
-        # Pretend that is_ returns a proper GelModel
-        @classmethod
-        def is_(
-            cls: _T_SelfModel, /, other_model: _T_OtherModel
-        ) -> type[GelModelIntersection[_T_SelfModel, _T_OtherModel]]: ...
-
-
-class GelModelIntersection(
-    GelModel, _abstract.BaseGelModelIntersection[_T_SelfModel, _T_OtherModel]
-):
-    def __init__(self) -> None:
-        raise NotImplementedError("Type expressions cannot be instantiated.")
-
 
 class GelLinkModel(
     GelSourceModel,

tests/test_qb.py

@@ -1581,7 +1581,6 @@ def test_qb_is_type_basic_03(self):
             excluded_fields={'b', 'c', 'ab', 'ac', 'bc', 'abc', 'ab_ac'},
         )
 
-    @tb.skip_typecheck
     def test_qb_is_type_basic_04(self):
         # Model Select
         # with computed single prop using type intersection