From 84eff755661c3c498c36e66b7ff0b6cc8fc3d4bb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sigurd=20Lj=C3=B8dal?= <544451+ljodal@users.noreply.github.com> Date: Thu, 30 Jun 2022 09:19:33 +0200 Subject: [PATCH] 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 * 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 * Add comment for test case * Fix typo * Rename variable Co-authored-by: Petter Friberg Co-authored-by: Nikita Sobolev --- mypy_django_plugin/transformers/managers.py | 116 +++++++++--------- .../managers/querysets/test_from_queryset.yml | 68 ++++++++-- 2 files changed, 113 insertions(+), 71 deletions(-) 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