From 59ebe6f1523c4d321b9e5ecba6c6458ace813489 Mon Sep 17 00:00:00 2001 From: Anthony Sottile Date: Fri, 26 Jul 2024 16:05:14 -0400 Subject: [PATCH] fill QuerySet generics using the manager's model type (#2281) --- mypy_django_plugin/main.py | 5 + mypy_django_plugin/transformers/managers.py | 57 +++++----- .../managers/querysets/test_as_manager.yml | 23 ++-- .../managers/querysets/test_from_queryset.yml | 100 ++++++++++++------ 4 files changed, 119 insertions(+), 66 deletions(-) diff --git a/mypy_django_plugin/main.py b/mypy_django_plugin/main.py index e9f549b70..554550378 100644 --- a/mypy_django_plugin/main.py +++ b/mypy_django_plugin/main.py @@ -37,6 +37,7 @@ ) from mypy_django_plugin.transformers.functional import resolve_str_promise_attribute from mypy_django_plugin.transformers.managers import ( + construct_as_manager_instance, create_new_manager_class_from_as_manager_method, create_new_manager_class_from_from_queryset_method, reparametrize_any_manager_hook, @@ -208,6 +209,10 @@ def get_method_hook(self, fullname: str) -> Optional[Callable[[MethodContext], M fullnames.REVERSE_MANY_TO_ONE_DESCRIPTOR: manytoone.refine_many_to_one_related_manager, } return hooks.get(class_fullname) + elif method_name == "as_manager": + info = self._get_typeinfo_or_none(class_fullname) + if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME): + return partial(construct_as_manager_instance, info=info) if method_name in self.manager_and_queryset_method_hooks: info = self._get_typeinfo_or_none(class_fullname) diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py index 951f1507b..ec13bf20c 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/mypy_django_plugin/transformers/managers.py @@ -15,11 +15,11 @@ StrExpr, SymbolTableNode, TypeInfo, - Var, ) -from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext +from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext, MethodContext from mypy.semanal import SemanticAnalyzer from mypy.semanal_shared import has_placeholder +from mypy.subtypes import find_member from mypy.types import ( AnyType, CallableType, @@ -28,6 +28,7 @@ Overloaded, ProperType, TypeOfAny, + TypeType, TypeVarType, UnionType, get_proper_type, @@ -121,15 +122,11 @@ def _process_dynamic_method( variables = method_type.variables ret_type = method_type.ret_type - if not is_fallback_queryset: - queryset_instance = 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 plugin-generated querysets. - queryset_instance = Instance(queryset_info, [manager_instance.args[0], manager_instance.args[0]]) + manager_model = find_member("model", manager_instance, manager_instance) + assert isinstance(manager_model, TypeType), manager_model + manager_model_type = manager_model.item + + queryset_instance = Instance(queryset_info, (manager_model_type,) * len(queryset_info.type_vars)) # 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 @@ -554,23 +551,9 @@ def create_new_manager_class_from_as_manager_method(ctx: DynamicClassDefContext) manager_name=manager_class_name, manager_base=manager_base, ) + queryset_info.metadata.setdefault("django_as_manager_names", {}) + queryset_info.metadata["django_as_manager_names"][semanal_api.cur_mod_id] = new_manager_info.name - # Whenever `.as_manager()` isn't called at class level, we want to ensure - # that the variable is an instance of our generated manager. Instead of the return - # value of `.as_manager()`. Though model argument is populated as `Any`. - # `transformers.models.AddManagers` will populate a model's manager(s), when it - # finds it on class level. - var = Var(name=ctx.name, type=Instance(new_manager_info, [AnyType(TypeOfAny.from_omitted_generics)])) - var.info = new_manager_info - var._fullname = f"{current_module.fullname}.{ctx.name}" - var.is_inferred = True - # Note: Order of `add_symbol_table_node` calls matters. Depending on what level - # we've found the `.as_manager()` call. Point here being that we want to replace the - # `.as_manager` return value with our newly created manager. - added = semanal_api.add_symbol_table_node( - ctx.name, SymbolTableNode(semanal_api.current_symbol_kind(), var, plugin_generated=True) - ) - assert added # Add the new manager to the current module added = semanal_api.add_symbol_table_node( # We'll use `new_manager_info.name` instead of `manager_class_name` here @@ -582,6 +565,26 @@ def create_new_manager_class_from_as_manager_method(ctx: DynamicClassDefContext) assert added +def construct_as_manager_instance(ctx: MethodContext, *, info: TypeInfo) -> MypyType: + api = helpers.get_typechecker_api(ctx) + module = helpers.get_current_module(api) + try: + manager_name = info.metadata["django_as_manager_names"][module.fullname] + except KeyError: + return ctx.default_return_type + + manager_node = api.lookup(manager_name) + if not isinstance(manager_node.node, TypeInfo): + return ctx.default_return_type + + # Whenever `.as_manager()` isn't called at class level, we want to ensure + # that the variable is an instance of our generated manager. Instead of the return + # value of `.as_manager()`. Though model argument is populated as `Any`. + # `transformers.models.AddManagers` will populate a model's manager(s), when it + # finds it on class level. + return Instance(manager_node.node, [AnyType(TypeOfAny.from_omitted_generics)]) + + def reparametrize_any_manager_hook(ctx: ClassDefContext) -> None: """ Add implicit generics to manager classes that are defined without generic. diff --git a/tests/typecheck/managers/querysets/test_as_manager.yml b/tests/typecheck/managers/querysets/test_as_manager.yml index 1c06b7a1a..641d4be30 100644 --- a/tests/typecheck/managers/querysets/test_as_manager.yml +++ b/tests/typecheck/managers/querysets/test_as_manager.yml @@ -14,13 +14,15 @@ - path: myapp/models.py content: | from django.db import models - from typing import List, Dict + from typing import List, Dict, TypeVar, ClassVar from typing_extensions import Self - class BaseQuerySet(models.QuerySet): + M = TypeVar("M", bound=models.Model, covariant=True) + + class BaseQuerySet(models.QuerySet[M]): def example_dict(self) -> Dict[str, Self]: ... - class MyQuerySet(BaseQuerySet): + class MyQuerySet(BaseQuerySet[M]): def example_simple(self) -> Self: ... def example_list(self) -> List[Self]: ... def just_int(self) -> int: ... @@ -64,9 +66,12 @@ - path: myapp/__init__.py - path: myapp/models.py content: | + from typing import TypeVar from django.db import models - class MyQuerySet(models.QuerySet): + M = TypeVar("M", bound=models.Model, covariant=True) + + class MyQuerySet(models.QuerySet[M]): ... class MyModel(models.Model): @@ -183,7 +188,7 @@ from myapp.models import MyModel, MyModelManager reveal_type(MyModelManager) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[Any]" reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]" - reveal_type(MyModel.objects.all()) # N: Revealed type is "myapp.models.ModelQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects.all()) # N: Revealed type is "myapp.models.ModelQuerySet" installed_apps: - myapp files: @@ -204,7 +209,7 @@ from myapp.models import MyModel, ManagerFromModelQuerySet reveal_type(ManagerFromModelQuerySet) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet1[Any]" reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet1[myapp.models.MyModel]" - reveal_type(MyModel.objects.all()) # N: Revealed type is "myapp.models.ModelQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects.all()) # N: Revealed type is "myapp.models.ModelQuerySet" installed_apps: - myapp files: @@ -280,7 +285,7 @@ objects = MyModelQuerySet.as_manager() class MyOtherModel(models.Model): - objects = _MyModelQuerySet2.as_manager() # type: ignore + objects = _MyModelQuerySet2.as_manager() - case: handles_type_vars main: | @@ -346,8 +351,8 @@ from myapp.models import MyModel reveal_type(MyModel.objects_1) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]" reveal_type(MyModel.objects_2) # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]" - reveal_type(MyModel.objects_1.all()) # N: Revealed type is "myapp.models.ModelQuerySet[myapp.models.MyModel]" - reveal_type(MyModel.objects_2.all()) # N: Revealed type is "myapp.models.ModelQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects_1.all()) # N: Revealed type is "myapp.models.ModelQuerySet" + reveal_type(MyModel.objects_2.all()) # N: Revealed type is "myapp.models.ModelQuerySet" installed_apps: - myapp files: diff --git a/tests/typecheck/managers/querysets/test_from_queryset.yml b/tests/typecheck/managers/querysets/test_from_queryset.yml index cb1da33b6..56f62d023 100644 --- a/tests/typecheck/managers/querysets/test_from_queryset.yml +++ b/tests/typecheck/managers/querysets/test_from_queryset.yml @@ -15,16 +15,18 @@ content: | from django.db import models from django.db.models.manager import BaseManager - from typing import List, Dict + from typing import List, Dict, TypeVar from typing_extensions import Self - class CustomManager(BaseManager): + M = TypeVar("M", covariant=True, bound=models.Model) + + class CustomManager(BaseManager[M]): def test_custom_manager(self) -> Self: ... - class BaseQuerySet(models.QuerySet): + class BaseQuerySet(models.QuerySet[M]): def example_dict(self) -> Dict[str, Self]: ... - class MyQuerySet(BaseQuerySet): + class MyQuerySet(BaseQuerySet[M]): def example_simple(self) -> Self: ... def example_list(self) -> List[Self]: ... def just_int(self) -> int: ... @@ -82,10 +84,13 @@ - path: myapp/__init__.py - path: myapp/models.py content: | + from typing import TypeVar from django.db import models from django.db.models.manager import BaseManager - class ModelQuerySet(models.QuerySet): + M = TypeVar("M", bound=models.Model, covariant=True) + + class ModelQuerySet(models.QuerySet[M]): def queryset_method(self) -> str: return 'hello' NewManager = BaseManager.from_queryset(ModelQuerySet) @@ -103,7 +108,7 @@ reveal_type(MyModel.objects.queryset_method_3()) # N: Revealed type is "builtins.str" reveal_type(MyModel.objects.queryset_method_4([])) # N: Revealed type is "None" reveal_type(MyModel.objects.filter(id=1).queryset_method()) # N: Revealed type is "myapp.querysets.ModelQuerySet" - reveal_type(MyModel.objects.filter(id=1)) # N: Revealed type is "myapp.querysets.ModelQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects.filter(id=1)) # N: Revealed type is "myapp.querysets.ModelQuerySet" installed_apps: - myapp files: @@ -223,7 +228,7 @@ reveal_type(MyModel.objects.queryset_method_3()) # N: Revealed type is "builtins.str" reveal_type(MyModel.objects.queryset_method_4([])) # N: Revealed type is "None" reveal_type(MyModel.objects.filter(id=1).queryset_method()) # N: Revealed type is "myapp.querysets.ModelQuerySet" - reveal_type(MyModel.objects.filter(id=1)) # N: Revealed type is "myapp.querysets.ModelQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects.filter(id=1)) # N: Revealed type is "myapp.querysets.ModelQuerySet" installed_apps: - myapp files: @@ -307,7 +312,7 @@ import typing kls: typing.Type[typing.Union[MyModel1, MyModel2]] = MyModel1 reveal_type(kls.objects) # N: Revealed type is "Union[myapp.models.ManagerFromModelQuerySet1[myapp.models.MyModel1], myapp.models.ManagerFromModelQuerySet2[myapp.models.MyModel2]]" - reveal_type(kls.objects.all()) # N: Revealed type is "Union[myapp.models.ModelQuerySet1[myapp.models.MyModel1], myapp.models.ModelQuerySet2[myapp.models.MyModel2]]" + reveal_type(kls.objects.all()) # N: Revealed type is "Union[myapp.models.ModelQuerySet1, myapp.models.ModelQuerySet2]" reveal_type(kls.objects.get()) # N: Revealed type is "Union[myapp.models.MyModel1, myapp.models.MyModel2]" reveal_type(kls.objects.queryset_method()) # N: Revealed type is "Union[builtins.int, builtins.str]" installed_apps: @@ -580,9 +585,9 @@ from myapp.models import MyModel reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]" reveal_type(MyModel._default_manager) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]" - reveal_type(MyModel.objects.all) # N: Revealed type is "def () -> myapp.models.MyQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects.all) # N: Revealed type is "def () -> myapp.models.MyQuerySet" reveal_type(MyModel.objects.custom) # N: Revealed type is "def () -> myapp.models.MyQuerySet" - reveal_type(MyModel.objects.all().filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet[myapp.models.MyModel]" + reveal_type(MyModel.objects.all().filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet" reveal_type(MyModel.objects.custom().filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet" reveal_type(MyModel.objects2) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]" reveal_type(MyModel._default_manager) # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]" @@ -633,26 +638,26 @@ - case: from_queryset_includes_methods_returning_queryset main: | from myapp.models import 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.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.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.distinct) # N: Revealed type is "def (*field_names: 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[typing.Sequence[builtins.str], None] =, params: Union[typing.Sequence[Any], None] =, tables: Union[typing.Sequence[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]" + reveal_type(MyModel.objects.alias) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.all) # N: Revealed type is "def () -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.annotate) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.complex_filter) # N: Revealed type is "def (filter_obj: Any) -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.defer) # N: Revealed type is "def (*fields: Any) -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.difference) # N: Revealed type is "def (*other_qs: Any) -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.distinct) # N: Revealed type is "def (*field_names: Any) -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.exclude) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.extra) # N: Revealed type is "def (select: Union[builtins.dict[builtins.str, Any], None] =, where: Union[typing.Sequence[builtins.str], None] =, params: Union[typing.Sequence[Any], None] =, tables: Union[typing.Sequence[builtins.str], None] =, order_by: Union[typing.Sequence[builtins.str], None] =, select_params: Union[typing.Sequence[Any], None] =) -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.intersection) # N: Revealed type is "def (*other_qs: Any) -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.none) # N: Revealed type is "def () -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.only) # N: Revealed type is "def (*fields: Any) -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.order_by) # N: Revealed type is "def (*field_names: Any) -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.prefetch_related) # N: Revealed type is "def (*lookups: Any) -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.reverse) # N: Revealed type is "def () -> myapp.models.MyQuerySet" + 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" + reveal_type(MyModel.objects.select_related) # N: Revealed type is "def (*fields: Any) -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.union) # N: Revealed type is "def (*other_qs: Any, all: builtins.bool =) -> myapp.models.MyQuerySet" + reveal_type(MyModel.objects.using) # N: Revealed type is "def (alias: Union[builtins.str, None]) -> myapp.models.MyQuerySet" installed_apps: - myapp files: @@ -882,6 +887,41 @@ class MCS(type): pass +- case: test_from_queryset_with_concrete_subclass + main: | + from myapp.models import Concrete + reveal_type(Concrete.objects) # N: Revealed type is "myapp.models.ConcreteManager" + reveal_type(Concrete.objects.get()) # N: Revealed type is "myapp.models.Concrete" + reveal_type(Concrete.objects.all()) # N: Revealed type is "myapp.models.CustomQuerySet[myapp.models.Concrete, myapp.models.Concrete]" + reveal_type(Concrete.objects.all().get()) # N: Revealed type is "myapp.models.Concrete" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from typing import ClassVar + from typing_extensions import Self, TypeVar + from django.db.models import Model, QuerySet + from django.db.models.manager import Manager + + M = TypeVar("M", bound=Model, covariant=True) + D = TypeVar("D", covariant=True, default=M) + + class CustomQuerySet(QuerySet[M, D]): pass + + _base = Manager.from_queryset(CustomQuerySet) + + class CustomBase(_base[M]): ... + + class BaseModel(Model): + objects: ClassVar[CustomBase[Self]] = CustomBase() + + class ConcreteManager(CustomBase["Concrete"]): ... + + class Concrete(BaseModel): + objects: ClassVar[ConcreteManager] = ConcreteManager() + - case: test_queryset_arg_as_unsupported_expressions main: | from typing import Union, Generic, TypeVar