Use bounded covariant TypeVars for explicit-arg QuerySet subclasses; restore master reparametrize_generic_class semantics for Manager/Field

Author: UnknownPlatypus

mypy_django_plugin/transformers/managers.py

@@ -606,10 +606,25 @@ def _defer() -> None:
     )
 
 
-def reparametrize_generic_class(ctx: ClassDefContext, base_class_fullname: str) -> None:
+def _is_omitted_generic(arg: MypyType) -> bool:
+    """True if ``arg`` is the ``Any`` mypy substitutes for an unparametrized generic."""
+    proper_arg = get_proper_type(arg)
+    return isinstance(proper_arg, AnyType) and proper_arg.type_of_any is TypeOfAny.from_omitted_generics
+
+
+def reparametrize_generic_class(
+    ctx: ClassDefContext,
+    base_class_fullname: str,
+    *,
+    bind_explicit_args: bool = False,
+) -> None:
     """
     Add implicit generics to classes that are defined without generic.
 
+    When ``bind_explicit_args`` is True, also reparametrize partially or
+    fully-explicit parametrizations by binding each explicit arg as both
+    the bound and the PEP 696 default of a synthesized covariant TypeVar.
+
     Note that this does not happen if mypy is run with disallow_any_generics = True,
     as not specifying the generic type is then considered an error.
     """
@@ -633,10 +648,31 @@ def reparametrize_generic_class(ctx: ClassDefContext, base_class_fullname: str)
     if parent_class is None or not parent_class.args:
         return
 
-    type_vars = [
-        type_var.copy_modified(id=type_var.id, default=parent_type_var)
-        for type_var, parent_type_var in zip(parent_class.type.defn.type_vars, parent_class.args, strict=True)
-    ]
+    # Bind explicit args only when the direct parent is the canonical base class
+    is_direct_parent = parent_class.type.fullname == base_class_fullname
+    if not (bind_explicit_args and is_direct_parent):
+        if not _is_omitted_generic(parent_class.args[0]):
+            return
+        type_vars = list(parent_class.type.defn.type_vars)
+    else:
+        type_vars = []
+        for type_var, parent_type_var in zip(parent_class.type.defn.type_vars, parent_class.args, strict=True):
+            if _is_omitted_generic(parent_type_var):
+                # Arg was omitted -- reuse the parent's TypeVar so its existing
+                # PEP 696 default (e.g. ``_Row = TypeVar(default=_Model)``) is preserved.
+                type_vars.append(type_var)
+            else:
+                # Arg was supplied explicitly -- synthesize a TypeVar bounded by the
+                # user's arg so the subclass stays effectively concrete in user-facing
+                # contexts (defaults + bound + covariance) while still being generic
+                # enough for the plugin to flow annotation row types through.
+                type_vars.append(
+                    type_var.copy_modified(
+                        id=type_var.id,
+                        upper_bound=parent_type_var,
+                        default=parent_type_var,
+                    )
+                )
 
     # If we end up with placeholders we need to defer so the placeholders are
     # resolved in a future iteration
@@ -654,6 +690,7 @@ def reparametrize_generic_class(ctx: ClassDefContext, base_class_fullname: str)
 def reparametrize_any_queryset_hook(ctx: ClassDefContext) -> None:
     """
     Add implicit generics to QuerySet subclasses that are defined without generic.
+    This is required for `.annotate` / `.values` call to update the typevars accordingly.
 
     Eg.
 
@@ -665,10 +702,20 @@ class MyQuerySet(models.QuerySet): ...
         _Row = TypeVar('_Row', covariant=True, default=_Model)
         class MyQuerySet(models.QuerySet[_Model, _Row]): ...
 
+    And
+
+        class BookQuerySet(models.QuerySet["Book"]): ...
+
+    is interpreted as:
+
+        _Model = TypeVar('_Model', bound=Book, covariant=True, default=Book)
+        _Row = TypeVar('_Row', bound=Book, covariant=True, default=Book)
+        class BookQuerySet(models.QuerySet[_Model, _Row]): ...
+
     Note that this does not happen if mypy is run with disallow_any_generics = True,
     as not specifying the generic type is then considered an error.
     """
-    reparametrize_generic_class(ctx, fullnames.QUERYSET_CLASS_FULLNAME)
+    reparametrize_generic_class(ctx, fullnames.QUERYSET_CLASS_FULLNAME, bind_explicit_args=True)
 
 
 def reparametrize_any_manager_hook(ctx: ClassDefContext) -> None:

tests/typecheck/fields/test_related.yml

@@ -906,15 +906,15 @@
         from myapp.models.store import Store
         from myapp.models.user import User
         reveal_type(Store().purchases)  # N: Revealed type is "myapp.models.purchase.Purchase_RelatedManager"
-        reveal_type(Store().purchases.queryset_method())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet"
-        reveal_type(Store().purchases.filter())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet"
+        reveal_type(Store().purchases.queryset_method())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet[myapp.models.purchase.Purchase, myapp.models.purchase.Purchase]"
+        reveal_type(Store().purchases.filter())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet[myapp.models.purchase.Purchase, myapp.models.purchase.Purchase]"
         reveal_type(Store().purchases.get())  # N: Revealed type is "myapp.models.purchase.Purchase"
-        reveal_type(Store().purchases.filter().queryset_method())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet"
+        reveal_type(Store().purchases.filter().queryset_method())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet[myapp.models.purchase.Purchase, myapp.models.purchase.Purchase]"
         reveal_type(User().purchases)  # N: Revealed type is "myapp.models.purchase.Purchase_RelatedManager"
-        reveal_type(User().purchases.queryset_method())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet"
-        reveal_type(User().purchases.filter())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet"
+        reveal_type(User().purchases.queryset_method())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet[myapp.models.purchase.Purchase, myapp.models.purchase.Purchase]"
+        reveal_type(User().purchases.filter())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet[myapp.models.purchase.Purchase, myapp.models.purchase.Purchase]"
         reveal_type(User().purchases.get())  # N: Revealed type is "myapp.models.purchase.Purchase"
-        reveal_type(User().purchases.filter().queryset_method())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet"
+        reveal_type(User().purchases.filter().queryset_method())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet[myapp.models.purchase.Purchase, myapp.models.purchase.Purchase]"
     installed_apps:
         - myapp
     files:

tests/typecheck/managers/querysets/test_as_manager.yml

@@ -6,7 +6,7 @@
         reveal_type(MyModel.objects.example_list())  # N: Revealed type is "list[myapp.models.MyQuerySet[myapp.models.MyModel]]"
         reveal_type(MyModel.objects.example_simple().just_int())  # N: Revealed type is "int"
         reveal_type(MyModel.objects.example_dict())  # N: Revealed type is "dict[str, myapp.models.MyQuerySet[myapp.models.MyModel]]"
-        reveal_type(MyModelWithoutSelf.objects.method())  # N: Revealed type is "myapp.models.QuerySetWithoutSelf"
+        reveal_type(MyModelWithoutSelf.objects.method())  # N: Revealed type is "myapp.models.QuerySetWithoutSelf[myapp.models.MyModelWithoutSelf, myapp.models.MyModelWithoutSelf]"
 
     installed_apps:
         - myapp
@@ -129,8 +129,8 @@
     main: |
         from typing_extensions import reveal_type
         from myapp.models import MyModel
-        reveal_type(MyModel.objects.custom_queryset_method())  # N: Revealed type is "myapp.models.ModelQuerySet"
-        reveal_type(MyModel.objects.all().custom_queryset_method())  # N: Revealed type is "myapp.models.ModelQuerySet"
+        reveal_type(MyModel.objects.custom_queryset_method())  # N: Revealed type is "myapp.models.ModelQuerySet[myapp.models.MyModel, myapp.models.MyModel]"
+        reveal_type(MyModel.objects.all().custom_queryset_method())  # N: Revealed type is "myapp.models.ModelQuerySet[myapp.models.MyModel, myapp.models.MyModel]"
         reveal_type(MyModel.objects.returns_int_sequence())  # N: Revealed type is "typing.Sequence[int]"
     installed_apps:
         - myapp
@@ -156,8 +156,8 @@
         from typing_extensions import reveal_type
         from myapp.models import MyModel1, MyModel2
         kls: type[MyModel1 | MyModel2] = MyModel1
-        reveal_type(kls.objects.custom_queryset_method())  # N: Revealed type is "myapp.models.ModelQuerySet1 | myapp.models.ModelQuerySet2"
-        reveal_type(kls.objects.all().custom_queryset_method())  # N: Revealed type is "myapp.models.ModelQuerySet1 | myapp.models.ModelQuerySet2"
+        reveal_type(kls.objects.custom_queryset_method())  # N: Revealed type is "myapp.models.ModelQuerySet1[myapp.models.MyModel1, myapp.models.MyModel1] | myapp.models.ModelQuerySet2[myapp.models.MyModel2, myapp.models.MyModel2]"
+        reveal_type(kls.objects.all().custom_queryset_method())  # N: Revealed type is "myapp.models.ModelQuerySet1[myapp.models.MyModel1, myapp.models.MyModel1] | myapp.models.ModelQuerySet2[myapp.models.MyModel2, myapp.models.MyModel2]"
         reveal_type(kls.objects.returns_thing())  # N: Revealed type is "int | str"
         reveal_type(kls.objects.get())  # N: Revealed type is "myapp.models.MyModel1 | myapp.models.MyModel2"
     installed_apps:
@@ -195,7 +195,7 @@
         from myapp.models import MyModel, MyModelManager
         reveal_type(MyModelManager)  # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[Any]"
         reveal_type(MyModel.objects)  # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]"
-        reveal_type(MyModel.objects.all())  # N: Revealed type is "myapp.models.ModelQuerySet"
+        reveal_type(MyModel.objects.all())  # N: Revealed type is "myapp.models.ModelQuerySet[myapp.models.MyModel, myapp.models.MyModel]"
     installed_apps:
         - myapp
     files:
@@ -217,7 +217,7 @@
         from myapp.models import MyModel, ManagerFromModelQuerySet
         reveal_type(ManagerFromModelQuerySet)  # N: Revealed type is "myapp.models.ManagerFromModelQuerySet1[Any]"
         reveal_type(MyModel.objects)  # N: Revealed type is "myapp.models.ManagerFromModelQuerySet1[myapp.models.MyModel]"
-        reveal_type(MyModel.objects.all())  # N: Revealed type is "myapp.models.ModelQuerySet"
+        reveal_type(MyModel.objects.all())  # N: Revealed type is "myapp.models.ModelQuerySet[myapp.models.MyModel, myapp.models.MyModel]"
     installed_apps:
         - myapp
     files:
@@ -371,8 +371,8 @@
         from myapp.models import MyModel
         reveal_type(MyModel.objects_1)  # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]"
         reveal_type(MyModel.objects_2)  # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]"
-        reveal_type(MyModel.objects_1.all())  # N: Revealed type is "myapp.models.ModelQuerySet"
-        reveal_type(MyModel.objects_2.all())  # N: Revealed type is "myapp.models.ModelQuerySet"
+        reveal_type(MyModel.objects_1.all())  # N: Revealed type is "myapp.models.ModelQuerySet[myapp.models.MyModel, myapp.models.MyModel]"
+        reveal_type(MyModel.objects_2.all())  # N: Revealed type is "myapp.models.ModelQuerySet[myapp.models.MyModel, myapp.models.MyModel]"
     installed_apps:
         - myapp
     files: