Fix row type of FallbackQuerySet

Because this inherits from _QuerySet, not QuerySet, it needs to have two
parameters

Author: ljodal

mypy_django_plugin/transformers/managers.py

@@ -87,11 +87,21 @@ def get_funcdef_type(definition: Union[FuncBase, Decorator, None]) -> Optional[P
     variables = method_type.variables
     ret_type = method_type.ret_type
 
+    is_fallback_queryset = queryset_info.metadata.get("django", {}).get("any_fallback_queryset", False)
+
     # For methods on the manager that return a queryset we need to override the
     # return type to be the actual queryset class, not the base QuerySet that's
     # used by the typing stubs.
     if method_name in MANAGER_METHODS_RETURNING_QUERYSET:
-        ret_type = Instance(queryset_info, manager_instance.args)
+        if not is_fallback_queryset:
+            ret_type = Instance(queryset_info, manager_instance.args)
+        else:
+            # The fallback queryset inherits _QuerySet, which has two generics
+            # instead of the one exposed on QuerySet. That means that we need
+            # to add the model twice. In real code it's not possible to inherit
+            # from _QuerySet, as it doesn't exist at runtime, so this fix is
+            # only needed for pluign-generated querysets.
+            ret_type = Instance(queryset_info, [manager_instance.args[0], manager_instance.args[0]])
         variables = []
 
     # Drop any 'self' argument as our manager is already initialized

mypy_django_plugin/transformers/models.py

@@ -160,8 +160,8 @@ def get_or_create_queryset_with_any_fallback(self) -> TypeInfo:
         }
         queryset_info.fallback_to_any = True
 
-        queryset_info.type_vars = base_queryset_info.type_vars
-        queryset_info.defn.type_vars = base_queryset_info.defn.type_vars
+        queryset_info.type_vars = base_queryset_info.type_vars.copy()
+        queryset_info.defn.type_vars = base_queryset_info.defn.type_vars.copy()
         queryset_info.metaclass_type = queryset_info.calculate_metaclass_type()
 
         return queryset_info

tests/typecheck/managers/test_managers.yml

@@ -508,12 +508,18 @@
                     reveal_type(Booking.objects.custom)
                     reveal_type(Booking.objects.all().filter)
                     reveal_type(Booking.objects.all().custom)
+                    reveal_type(Booking.objects.first())
+                    reveal_type(Booking.objects.get())
+                    reveal_type([booking for booking in Booking.objects.all()])
+                    reveal_type([booking for booking in Booking.objects.all().filter()])
+
 
                     # Check QuerySet methods on UnknownRelatedManager
                     reveal_type(user.booking_set.all)
                     reveal_type(user.booking_set.custom)
                     reveal_type(user.booking_set.all().filter)
                     reveal_type(user.booking_set.all().custom)
+                    reveal_type(user.booking_set.all().first())
     out: |
         myapp/models:13: error: Couldn't resolve related manager for relation 'booking' (from myapp.models.Booking.myapp.Booking.renter).
         myapp/models:13: error: Couldn't resolve related manager for relation 'bookingowner_set' (from myapp.models.Booking.myapp.Booking.owner).
@@ -533,11 +539,16 @@
         myapp/models:47: note: Revealed type is "django.db.models.manager.BaseManager[myapp.models.InvisibleUnresolvable]"
         myapp/models:49: note: Revealed type is "myapp.models.UnknownRelatedManager[myapp.models.Booking]"
         myapp/models:50: note: Revealed type is "myapp.models.UnknownRelatedManager[myapp.models.Booking]"
-        myapp/models:53: note: Revealed type is "def () -> myapp.models.UnknownQuerySet[myapp.models.Booking]"
+        myapp/models:53: note: Revealed type is "def () -> myapp.models.UnknownQuerySet[myapp.models.Booking, myapp.models.Booking]"
         myapp/models:54: note: Revealed type is "Any"
-        myapp/models:55: note: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.UnknownQuerySet[myapp.models.Booking]"
+        myapp/models:55: note: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.UnknownQuerySet[myapp.models.Booking, myapp.models.Booking]"
         myapp/models:56: note: Revealed type is "Any"
-        myapp/models:59: note: Revealed type is "def () -> myapp.models.UnknownQuerySet[myapp.models.Booking]"
-        myapp/models:60: note: Revealed type is "Any"
-        myapp/models:61: note: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.UnknownQuerySet[myapp.models.Booking]"
-        myapp/models:62: note: Revealed type is "Any"
+        myapp/models:57: note: Revealed type is "Union[myapp.models.Booking, None]"
+        myapp/models:58: note: Revealed type is "myapp.models.Booking"
+        myapp/models:59: note: Revealed type is "builtins.list[myapp.models.Booking]"
+        myapp/models:60: note: Revealed type is "builtins.list[myapp.models.Booking]"
+        myapp/models:64: note: Revealed type is "def () -> myapp.models.UnknownQuerySet[myapp.models.Booking, myapp.models.Booking]"
+        myapp/models:65: note: Revealed type is "Any"
+        myapp/models:66: note: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.UnknownQuerySet[myapp.models.Booking, myapp.models.Booking]"
+        myapp/models:67: note: Revealed type is "Any"
+        myapp/models:68: note: Revealed type is "Union[myapp.models.Booking, None]"