Improve typing for unresolved managers (#1044)

* Improve typing for unresolved managers

This changes the logic when encountering an unresolvable manager class.
Instead of adding it as a `Manager` we create a subclass of `Manager`
that has `fallback_to_any=True` set. Similarly a `QuerySet` class is
created that also has fallbacks to `Any`. This allows calling custom
methods on the manager and querysets without getting type errors.

* Fix manager created and improve a test

* Fix row type of FallbackQuerySet

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

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

@@ -12,14 +12,18 @@
 from mypy.types import AnyType, Instance
 from mypy.types import Type as MypyType
 from mypy.types import TypedDictType, TypeOfAny
+from mypy.typevars import fill_typevars
 
 from mypy_django_plugin.django.context import DjangoContext
 from mypy_django_plugin.errorcodes import MANAGER_MISSING
 from mypy_django_plugin.lib import fullnames, helpers
 from mypy_django_plugin.lib.fullnames import ANNOTATIONS_FULLNAME, ANY_ATTR_ALLOWED_CLASS_FULLNAME, MODEL_CLASS_FULLNAME
 from mypy_django_plugin.transformers import fields
 from mypy_django_plugin.transformers.fields import get_field_descriptor_types
-from mypy_django_plugin.transformers.managers import create_manager_info_from_from_queryset_call
+from mypy_django_plugin.transformers.managers import (
+    MANAGER_METHODS_RETURNING_QUERYSET,
+    create_manager_info_from_from_queryset_call,
+)
 
 
 class ModelClassInitializer:
@@ -83,6 +87,85 @@ def get_generated_manager_info(self, manager_fullname: str, base_manager_fullnam
         # Not a generated manager
         return None
 
+    def get_or_create_manager_with_any_fallback(self, related_manager: bool = False) -> TypeInfo:
+        """
+        Create a Manager subclass with fallback to Any for unknown attributes
+        and methods. This is used for unresolved managers, where we don't know
+        the actual type of the manager.
+
+        The created class is reused if multiple unknown managers are encountered.
+        """
+
+        name = "UnknownManager" if not related_manager else "UnknownRelatedManager"
+
+        # Check if we've already created a fallback manager class for this
+        # module, and if so reuse that.
+        manager_info = self.lookup_typeinfo(f"{self.model_classdef.info.module_name}.{name}")
+        if manager_info and manager_info.metadata.get("django", {}).get("any_fallback_manager"):
+            return manager_info
+
+        fallback_queryset = self.get_or_create_queryset_with_any_fallback()
+        base_manager_fullname = (
+            fullnames.MANAGER_CLASS_FULLNAME if not related_manager else fullnames.RELATED_MANAGER_CLASS
+        )
+        base_manager_info = self.lookup_typeinfo(base_manager_fullname)
+        assert base_manager_info, f"Type info for {base_manager_fullname} missing"
+
+        base_manager = fill_typevars(base_manager_info)
+        assert isinstance(base_manager, Instance)
+        manager_info = self.add_new_class_for_current_module(name, [base_manager])
+        manager_info.fallback_to_any = True
+
+        manager_info.type_vars = base_manager_info.type_vars
+        manager_info.defn.type_vars = base_manager_info.defn.type_vars
+        manager_info.metaclass_type = manager_info.calculate_metaclass_type()
+
+        # For methods on BaseManager that return a queryset we need to update
+        # the return type to be the actual queryset subclass used. This is done
+        # by adding the methods as attributes with type Any to the manager
+        # class. The actual type of these methods are resolved in
+        # resolve_manager_method.
+        for method_name in MANAGER_METHODS_RETURNING_QUERYSET:
+            helpers.add_new_sym_for_info(manager_info, name=method_name, sym_type=AnyType(TypeOfAny.special_form))
+
+        manager_info.metadata["django"] = django_metadata = {
+            "any_fallback_manager": True,
+            "from_queryset_manager": fallback_queryset.fullname,
+        }
+
+        return manager_info
+
+    def get_or_create_queryset_with_any_fallback(self) -> TypeInfo:
+        """
+        Create a QuerySet subclass with fallback to Any for unknown attributes
+        and methods. This is used for the manager returned by the method above.
+        """
+
+        name = "UnknownQuerySet"
+
+        # Check if we've already created a fallback queryset class for this
+        # module, and if so reuse that.
+        queryset_info = self.lookup_typeinfo(f"{self.model_classdef.info.module_name}.{name}")
+        if queryset_info and queryset_info.metadata.get("django", {}).get("any_fallback_queryset"):
+            return queryset_info
+
+        base_queryset_info = self.lookup_typeinfo(fullnames.QUERYSET_CLASS_FULLNAME)
+        assert base_queryset_info, f"Type info for {fullnames.QUERYSET_CLASS_FULLNAME} missing"
+
+        base_queryset = fill_typevars(base_queryset_info)
+        assert isinstance(base_queryset, Instance)
+        queryset_info = self.add_new_class_for_current_module(name, [base_queryset])
+        queryset_info.metadata["django"] = {
+            "any_fallback_queryset": True,
+        }
+        queryset_info.fallback_to_any = True
+
+        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
+
     def run_with_model_cls(self, model_cls):
         raise NotImplementedError("Implement this in subclasses")
 
@@ -285,13 +368,11 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
                 # (Django's default) before finishing. And emit an error, to allow for
                 # ignoring a more specialised manager not being resolved while still
                 # setting _some_ type
-                django_manager_info = self.lookup_typeinfo(fullnames.MANAGER_CLASS_FULLNAME)
-                assert (
-                    django_manager_info is not None
-                ), f"Type info for Django's {fullnames.MANAGER_CLASS_FULLNAME} missing"
+                fallback_manager_info = self.get_or_create_manager_with_any_fallback()
                 self.add_new_node_to_model_class(
-                    manager_name, Instance(django_manager_info, [Instance(self.model_classdef.info, [])])
+                    manager_name, Instance(fallback_manager_info, [Instance(self.model_classdef.info, [])])
                 )
+
                 # Find expression for e.g. `objects = SomeManager()`
                 manager_expr = self.get_manager_expression(manager_name)
                 manager_fullname = f"{self.model_classdef.fullname}.{manager_name}"
@@ -425,28 +506,27 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
                 except helpers.IncompleteDefnException as exc:
                     if not self.api.final_iteration:
                         raise exc
-                    else:
-                        if related_manager_info:
-                            """
-                            If a django model has a Manager class that cannot be
-                            resolved statically (if it is generated in a way
-                            where we cannot import it, like `objects = my_manager_factory()`),
-                            we fallback to the default related manager, so you
-                            at least get a base level of working type checking.
-
-                            See https://github.com/typeddjango/django-stubs/pull/993
-                            for more information on when this error can occur.
-                            """
-                            self.add_new_node_to_model_class(
-                                attname, Instance(related_manager_info, [Instance(related_model_info, [])])
-                            )
-                            related_model_fullname = related_model_cls.__module__ + "." + related_model_cls.__name__
-                            self.ctx.api.fail(
-                                f"Couldn't resolve related manager for relation {relation.name!r} (from {related_model_fullname}.{relation.field}).",
-                                self.ctx.cls,
-                                code=MANAGER_MISSING,
-                            )
-                        continue
+
+                    # If a django model has a Manager class that cannot be
+                    # resolved statically (if it is generated in a way where we
+                    # cannot import it, like `objects = my_manager_factory()`),
+                    # we fallback to the default related manager, so you at
+                    # least get a base level of working type checking.
+                    #
+                    # See https://github.com/typeddjango/django-stubs/pull/993
+                    # for more information on when this error can occur.
+                    fallback_manager = self.get_or_create_manager_with_any_fallback(related_manager=True)
+                    self.add_new_node_to_model_class(
+                        attname, Instance(fallback_manager, [Instance(related_model_info, [])])
+                    )
+                    related_model_fullname = related_model_cls.__module__ + "." + related_model_cls.__name__
+                    self.ctx.api.fail(
+                        f"Couldn't resolve related manager for relation {relation.name!r} (from {related_model_fullname}.{relation.field}).",
+                        self.ctx.cls,
+                        code=MANAGER_MISSING,
+                    )
+
+                    continue
 
                 # Check if the related model has a related manager subclassed from the default manager
                 # TODO: Support other reverse managers than `_default_manager`

tests/typecheck/fields/test_related.yml

@@ -678,18 +678,24 @@
                     def custom(self) -> None:
                         pass
 
+                def TransactionManager():
+                    return BaseManager.from_queryset(TransactionQuerySet)()
+
                 class Transaction(models.Model):
-                    objects = BaseManager.from_queryset(TransactionQuerySet)
+                    objects = TransactionManager()
                     def test(self) -> None:
                         reveal_type(self.transactionlog_set)
                         # We use a fallback Any type:
+                        reveal_type(Transaction.objects)
                         reveal_type(Transaction.objects.custom())
 
                 class TransactionLog(models.Model):
                     transaction = models.ForeignKey(Transaction, on_delete=models.CASCADE)
     out: |
-        myapp/models:10: note: Revealed type is "django.db.models.manager.RelatedManager[myapp.models.TransactionLog]"
-        myapp/models:12: note: Revealed type is "Any"
+        myapp/models:11: error: Could not resolve manager type for "myapp.models.Transaction.objects"
+        myapp/models:13: note: Revealed type is "django.db.models.manager.RelatedManager[myapp.models.TransactionLog]"
+        myapp/models:15: note: Revealed type is "myapp.models.UnknownManager[myapp.models.Transaction]"
+        myapp/models:16: note: Revealed type is "Any"
 
 
 -   case: resolve_primary_keys_for_foreign_keys_with_abstract_self_model

tests/typecheck/managers/test_managers.yml

@@ -502,6 +502,24 @@
 
                     reveal_type(user.bookingowner_set)
                     reveal_type(user.booking_set)
+
+                    # Check QuerySet methods on UnknownManager
+                    reveal_type(Booking.objects.all)
+                    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).
@@ -512,12 +530,25 @@
         myapp/models:32: error: Could not resolve manager type for "myapp.models.InvisibleUnresolvable.objects"
         myapp/models:36: note: Revealed type is "django.db.models.manager.Manager[myapp.models.User]"
         myapp/models:37: note: Revealed type is "django.db.models.manager.Manager[myapp.models.User]"
-        myapp/models:39: note: Revealed type is "django.db.models.manager.Manager[myapp.models.Booking]"
+        myapp/models:39: note: Revealed type is "myapp.models.UnknownManager[myapp.models.Booking]"
         myapp/models:40: note: Revealed type is "django.db.models.manager.BaseManager[myapp.models.Booking]"
-        myapp/models:42: note: Revealed type is "django.db.models.manager.Manager[myapp.models.TwoUnresolvable]"
-        myapp/models:43: note: Revealed type is "django.db.models.manager.Manager[myapp.models.TwoUnresolvable]"
+        myapp/models:42: note: Revealed type is "myapp.models.UnknownManager[myapp.models.TwoUnresolvable]"
+        myapp/models:43: note: Revealed type is "myapp.models.UnknownManager[myapp.models.TwoUnresolvable]"
         myapp/models:44: note: Revealed type is "django.db.models.manager.BaseManager[myapp.models.TwoUnresolvable]"
-        myapp/models:46: note: Revealed type is "django.db.models.manager.Manager[myapp.models.InvisibleUnresolvable]"
+        myapp/models:46: note: Revealed type is "myapp.models.UnknownManager[myapp.models.InvisibleUnresolvable]"
         myapp/models:47: note: Revealed type is "django.db.models.manager.BaseManager[myapp.models.InvisibleUnresolvable]"
-        myapp/models:49: note: Revealed type is "django.db.models.manager.RelatedManager[myapp.models.Booking]"
-        myapp/models:50: note: Revealed type is "django.db.models.manager.RelatedManager[myapp.models.Booking]"
+        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.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.Booking]"
+        myapp/models:56: 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]"