diff --git a/mypy_django_plugin/main.py b/mypy_django_plugin/main.py index ec89f0663..b724d7147 100644 --- a/mypy_django_plugin/main.py +++ b/mypy_django_plugin/main.py @@ -27,6 +27,7 @@ from mypy_django_plugin.transformers import fields, forms, init_create, manytomany, meta, querysets, request, settings 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, @@ -195,6 +196,10 @@ def get_method_hook(self, fullname: str) -> Optional[Callable[[MethodContext], M fullnames.MANY_TO_MANY_DESCRIPTOR, }: return manytomany.refine_many_to_many_related_manager + 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) manager_classes = self._get_current_manager_bases() diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py index 2da8e50de..6e49b0460 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/mypy_django_plugin/transformers/managers.py @@ -14,11 +14,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, @@ -27,6 +27,7 @@ Overloaded, ProperType, TypeOfAny, + TypeType, TypeVarType, UnionType, get_proper_type, @@ -120,15 +121,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 @@ -553,26 +550,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 - # So that the plugin will reparameterize the manager when it is constructed inside of a Model definition - helpers.add_new_manager_base(semanal_api, new_manager_info.fullname) - - # 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 @@ -584,6 +564,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: