Resolve all queryset methods on managers as attributes

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