diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py index 836d0e93a..28ac77418 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/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,7 +132,7 @@ 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 @@ -98,9 +140,10 @@ def get_method_type_from_reverse_manager( 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)) diff --git a/tests/typecheck/managers/querysets/test_from_queryset.yml b/tests/typecheck/managers/querysets/test_from_queryset.yml index caa8249d4..142f1f25a 100644 --- a/tests/typecheck/managers/querysets/test_from_queryset.yml +++ b/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