Resolve all queryset methods on managers as attributes (#1028)

* Add test case reproducing Sequence name not defined issue

* Resolve all manager methods as attribute

This changes to logic for resolving methods from the base QuerySet class
on managers from copying the methods to use the attribute approach
that's already used for methods from custom querysets. This resolves the
phantom type errors that stem from the copying.

* Disable cache in test case

Make sure the test will fail regardless of which mypy.ini file is being using.

Co-authored-by: Petter Friberg <[email protected]>

* Update comments related to copying methods

* Use a predefined list of manager methods to update

The list of manager methods that returns a queryset, and thus need to
have it's return type changed, is small and well defined. Using a
predefined list of methods rather than trying to detect these at runtime
makes the code much more readable and probably faster as well.

Also add `extra()` to the methods tested in
from_queryset_includes_methods_returning_queryset, and sort the methods
alphabetically.

* Revert changes in .github/workflows/tests.yml

With cache_disable: true on the test case this is no longer needed to
reproduce the bug.

* Remove unsued imports and change type of constant

 - Remove unused imports left behind
 - Change MANAGER_METHODS_RETURNING_QUERYSET to Final[FrozenSet[str]]

* Import Final from typing_extensions

Was added in 3.8, we still support 3.7

* Sort imports properly

* Remove explicit typing of final frozenset

Co-authored-by: Nikita Sobolev <[email protected]>

* Add comment for test case

* Fix typo

* Rename variable

Co-authored-by: Petter Friberg <[email protected]>
Co-authored-by: Nikita Sobolev <[email protected]>

mypy_django_plugin/transformers/managers.py

@@ -19,18 +19,47 @@
 from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext, MethodContext
 from mypy.types import AnyType, CallableType, Instance, ProperType
 from mypy.types import Type as MypyType
-from mypy.types import TypeOfAny, TypeVarType, UnboundType, get_proper_type
+from mypy.types import TypeOfAny
+from typing_extensions import Final
 
 from mypy_django_plugin import errorcodes
 from mypy_django_plugin.lib import fullnames, helpers
 
+MANAGER_METHODS_RETURNING_QUERYSET: Final = frozenset(
+    (
+        "alias",
+        "all",
+        "annotate",
+        "complex_filter",
+        "defer",
+        "difference",
+        "distinct",
+        "exclude",
+        "extra",
+        "filter",
+        "intersection",
+        "none",
+        "only",
+        "order_by",
+        "prefetch_related",
+        "reverse",
+        "select_for_update",
+        "select_related",
+        "union",
+        "using",
+    )
+)
+
 
 def get_method_type_from_dynamic_manager(
-    api: TypeChecker, method_name: str, manager_type_info: TypeInfo
+    api: TypeChecker, method_name: str, manager_instance: Instance
 ) -> Optional[ProperType]:
     """
     Attempt to resolve a method on a manager that was built from '.from_queryset'
     """
+
+    manager_type_info = manager_instance.type
+
     if (
         "django" not in manager_type_info.metadata
         or "from_queryset_manager" not in manager_type_info.metadata["django"]
@@ -56,11 +85,24 @@ def get_funcdef_type(definition: Union[FuncBase, Decorator, None]) -> Optional[P
         return None
 
     assert isinstance(method_type, CallableType)
+
+    variables = method_type.variables
+    ret_type = method_type.ret_type
+
+    # 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)
+        variables = []
+
     # Drop any 'self' argument as our manager is already initialized
     return method_type.copy_modified(
         arg_types=method_type.arg_types[1:],
         arg_kinds=method_type.arg_kinds[1:],
         arg_names=method_type.arg_names[1:],
+        variables=variables,
+        ret_type=ret_type,
     )
 
 
@@ -90,17 +132,18 @@ def get_method_type_from_reverse_manager(
     assert isinstance(model_info.names["_default_manager"].node, Var)
     manager_instance = model_info.names["_default_manager"].node.type
     return (
-        get_method_type_from_dynamic_manager(api, method_name, manager_instance.type)
+        get_method_type_from_dynamic_manager(api, method_name, manager_instance)
         # TODO: Can we assert on None and Instance?
         if manager_instance is not None and isinstance(manager_instance, Instance)
         else None
     )
 
 
 def resolve_manager_method_from_instance(instance: Instance, method_name: str, ctx: AttributeContext) -> MypyType:
+
     api = helpers.get_typechecker_api(ctx)
     method_type = get_method_type_from_dynamic_manager(
-        api, method_name, instance.type
+        api, method_name, instance
     ) or get_method_type_from_reverse_manager(api, method_name, instance.type)
 
     return method_type if method_type is not None else ctx.default_attr_type
@@ -232,60 +275,17 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
                 sym_type=AnyType(TypeOfAny.special_form),
             )
 
-    # we need to copy all methods in MRO before django.db.models.query.QuerySet
-    # Gather names of all BaseManager methods
-    manager_method_names = []
-    for manager_mro_info in new_manager_info.mro:
-        if manager_mro_info.fullname == fullnames.BASE_MANAGER_CLASS_FULLNAME:
-            for name, sym in manager_mro_info.names.items():
-                manager_method_names.append(name)
-
-    # Copy/alter all methods in common between BaseManager/QuerySet over to the new manager if their return type is
-    # the QuerySet's self-type. Alter the return type to be the custom queryset, parameterized by the manager's model
-    # type variable.
-    for class_mro_info in derived_queryset_info.mro:
-        if class_mro_info.fullname != fullnames.QUERYSET_CLASS_FULLNAME:
-            continue
-        for name, sym in class_mro_info.names.items():
-            if name not in manager_method_names:
-                continue
-
-            if isinstance(sym.node, FuncDef):
-                func_node = sym.node
-            elif isinstance(sym.node, Decorator):
-                func_node = sym.node.func
-            else:
-                continue
-
-            method_type = func_node.type
-            if not isinstance(method_type, CallableType):
-                if not semanal_api.final_iteration:
-                    semanal_api.defer()
-                return None
-            original_return_type = method_type.ret_type
-
-            # Skip any method that doesn't return _QS
-            original_return_type = get_proper_type(original_return_type)
-            if isinstance(original_return_type, UnboundType):
-                if original_return_type.name != "_QS":
-                    continue
-            elif isinstance(original_return_type, TypeVarType):
-                if original_return_type.name != "_QS":
-                    continue
-            else:
-                continue
-
-            # Return the custom queryset parameterized by the manager's type vars
-            return_type = Instance(derived_queryset_info, self_type.args)
-
-            helpers.copy_method_to_another_class(
-                class_def_context,
-                self_type,
-                new_method_name=name,
-                method_node=func_node,
-                return_type=return_type,
-                original_module_name=class_mro_info.module_name,
-            )
+    # 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,
+    # similar to how custom queryset methods are handled above. The actual type
+    # of these methods are resolved in resolve_manager_method.
+    for name in MANAGER_METHODS_RETURNING_QUERYSET:
+        helpers.add_new_sym_for_info(
+            new_manager_info,
+            name=name,
+            sym_type=AnyType(TypeOfAny.special_form),
+        )
 
     # Insert the new manager (dynamic) class
     assert semanal_api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, new_manager_info, plugin_generated=True))

tests/typecheck/managers/querysets/test_from_queryset.yml

@@ -350,24 +350,25 @@
 -   case: from_queryset_includes_methods_returning_queryset
     main: |
         from myapp.models import MyModel
-        reveal_type(MyModel.objects.none)  # N: Revealed type is "def () -> myapp.models.MyQuerySet[myapp.models.MyModel]"
+        reveal_type(MyModel.objects.alias)  # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
         reveal_type(MyModel.objects.all)  # N: Revealed type is "def () -> myapp.models.MyQuerySet[myapp.models.MyModel]"
-        reveal_type(MyModel.objects.filter)  # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
-        reveal_type(MyModel.objects.exclude)  # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
+        reveal_type(MyModel.objects.annotate)  # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
         reveal_type(MyModel.objects.complex_filter)  # N: Revealed type is "def (filter_obj: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
-        reveal_type(MyModel.objects.union)  # N: Revealed type is "def (*other_qs: Any, *, all: builtins.bool =) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
-        reveal_type(MyModel.objects.intersection)  # N: Revealed type is "def (*other_qs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
+        reveal_type(MyModel.objects.defer)  # N: Revealed type is "def (*fields: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
         reveal_type(MyModel.objects.difference)  # N: Revealed type is "def (*other_qs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
-        reveal_type(MyModel.objects.select_for_update)  # N: Revealed type is "def (nowait: builtins.bool =, skip_locked: builtins.bool =, of: typing.Sequence[builtins.str] =, no_key: builtins.bool =) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
-        reveal_type(MyModel.objects.select_related)  # N: Revealed type is "def (*fields: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
-        reveal_type(MyModel.objects.prefetch_related)  # N: Revealed type is "def (*lookups: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
-        reveal_type(MyModel.objects.annotate)  # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
-        reveal_type(MyModel.objects.alias)  # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
-        reveal_type(MyModel.objects.order_by)  # N: Revealed type is "def (*field_names: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
         reveal_type(MyModel.objects.distinct)  # N: Revealed type is "def (*field_names: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
-        reveal_type(MyModel.objects.reverse)  # N: Revealed type is "def () -> myapp.models.MyQuerySet[myapp.models.MyModel]"
-        reveal_type(MyModel.objects.defer)  # N: Revealed type is "def (*fields: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
+        reveal_type(MyModel.objects.exclude)  # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
+        reveal_type(MyModel.objects.extra)  # N: Revealed type is "def (select: Union[builtins.dict[builtins.str, Any], None] =, where: Union[builtins.list[builtins.str], None] =, params: Union[builtins.list[Any], None] =, tables: Union[builtins.list[builtins.str], None] =, order_by: Union[typing.Sequence[builtins.str], None] =, select_params: Union[typing.Sequence[Any], None] =) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
+        reveal_type(MyModel.objects.filter)  # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
+        reveal_type(MyModel.objects.intersection)  # N: Revealed type is "def (*other_qs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
+        reveal_type(MyModel.objects.none)  # N: Revealed type is "def () -> myapp.models.MyQuerySet[myapp.models.MyModel]"
         reveal_type(MyModel.objects.only)  # N: Revealed type is "def (*fields: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
+        reveal_type(MyModel.objects.order_by)  # N: Revealed type is "def (*field_names: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
+        reveal_type(MyModel.objects.prefetch_related)  # N: Revealed type is "def (*lookups: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
+        reveal_type(MyModel.objects.reverse)  # N: Revealed type is "def () -> myapp.models.MyQuerySet[myapp.models.MyModel]"
+        reveal_type(MyModel.objects.select_for_update)  # N: Revealed type is "def (nowait: builtins.bool =, skip_locked: builtins.bool =, of: typing.Sequence[builtins.str] =, no_key: builtins.bool =) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
+        reveal_type(MyModel.objects.select_related)  # N: Revealed type is "def (*fields: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
+        reveal_type(MyModel.objects.union)  # N: Revealed type is "def (*other_qs: Any, *, all: builtins.bool =) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
         reveal_type(MyModel.objects.using)  # N: Revealed type is "def (alias: Union[builtins.str, None]) -> myapp.models.MyQuerySet[myapp.models.MyModel]"
     installed_apps:
         - myapp
@@ -384,3 +385,44 @@
                 MyManager = BaseManager.from_queryset(MyQuerySet)
                 class MyModel(models.Model):
                     objects = MyManager()
+
+
+# This tests a regression where mypy would generate phantom warnings about
+# undefined types due to unresolved types when copying methods from QuerySet to
+# a manager dynamically created using Manager.from_queryset().
+#
+# For details see: https://github.com/typeddjango/django-stubs/issues/1022
+-   case: from_queryset_custom_auth_user_model
+    # Cache needs to be disabled to consistenly reproduce the bug
+    disable_cache: true
+    main: |
+        from users.models import User
+    custom_settings: |
+        AUTH_USER_MODEL = "users.User"
+        INSTALLED_APPS = ("django.contrib.auth", "django.contrib.contenttypes", "users")
+    files:
+        -   path: users/__init__.py
+        -   path: users/models.py
+            content: |
+                from django.contrib.auth.models import AbstractBaseUser
+                from django.db import models
+
+                from .querysets import UserQuerySet
+
+                UserManager = models.Manager.from_queryset(UserQuerySet)
+
+                class User(AbstractBaseUser):
+                    email = models.EmailField(unique=True)
+                    objects = UserManager()
+                    USERNAME_FIELD = "email"
+
+        -   path: users/querysets.py
+            content: |
+                from django.db import models
+                from typing import Optional, TYPE_CHECKING
+
+                if TYPE_CHECKING:
+                    from .models import User
+
+                class UserQuerySet(models.QuerySet["User"]):
+                    pass