From a0d8d6dfea07758ad9296e66700b216cc33016d5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sigurd=20Lj=C3=B8dal?= <544451+ljodal@users.noreply.github.com> Date: Sun, 3 Jul 2022 22:54:15 +0200 Subject: [PATCH] Always create manager at module level When the manager is added at the class level, which happened when it was created inline in the model body, it's not possible to retrieve the manager again based on fullname. That lead to problems with inheritance and the default manager. --- mypy_django_plugin/lib/helpers.py | 28 +++++++++++------ mypy_django_plugin/transformers/managers.py | 14 ++++----- mypy_django_plugin/transformers/models.py | 16 ++-------- .../managers/querysets/test_from_queryset.yml | 31 +++++++++++++++++-- 4 files changed, 56 insertions(+), 33 deletions(-) diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py index da19db01f1..a101c07f3a 100644 --- a/mypy_django_plugin/lib/helpers.py +++ b/mypy_django_plugin/lib/helpers.py @@ -208,6 +208,23 @@ def is_annotated_model_fullname(model_cls_fullname: str) -> bool: return model_cls_fullname.startswith(WITH_ANNOTATIONS_FULLNAME + "[") +def create_type_info(name: str, module: str, bases: list[Instance]) -> TypeInfo: + + # make new class expression + classdef = ClassDef(name, Block([])) + classdef.fullname = module + "." + name + + # make new TypeInfo + new_typeinfo = TypeInfo(SymbolTable(), classdef, module) + new_typeinfo.bases = bases + calculate_mro(new_typeinfo) + new_typeinfo.calculate_metaclass_type() + + classdef.info = new_typeinfo + + return new_typeinfo + + def add_new_class_for_module( module: MypyFile, name: str, @@ -217,15 +234,7 @@ def add_new_class_for_module( ) -> TypeInfo: new_class_unique_name = checker.gen_unique_name(name, module.names) - # make new class expression - classdef = ClassDef(new_class_unique_name, Block([])) - classdef.fullname = module.fullname + "." + new_class_unique_name - - # make new TypeInfo - new_typeinfo = TypeInfo(SymbolTable(), classdef, module.fullname) - new_typeinfo.bases = bases - calculate_mro(new_typeinfo) - new_typeinfo.calculate_metaclass_type() + new_typeinfo = create_type_info(new_class_unique_name, module.fullname, bases) # add fields if fields: @@ -237,7 +246,6 @@ def add_new_class_for_module( MDEF, var, plugin_generated=True, no_serialize=no_serialize ) - classdef.info = new_typeinfo module.names[new_class_unique_name] = SymbolTableNode( GDEF, new_typeinfo, plugin_generated=True, no_serialize=no_serialize ) diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py index 6431524f22..ce41b6ee99 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/mypy_django_plugin/transformers/managers.py @@ -199,12 +199,6 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte # 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) - # Insert the new manager (dynamic) class - assert semanal_api.add_symbol_table_node( - ctx.name, - SymbolTableNode(GDEF, new_manager_info, plugin_generated=True), - ) - def create_manager_info_from_from_queryset_call( api: SemanticAnalyzerPluginInterface, call_expr: CallExpr, name: Optional[str] = None @@ -251,6 +245,12 @@ def create_manager_info_from_from_queryset_call( base_manager_info.metadata.setdefault("from_queryset_managers", {}) base_manager_info.metadata["from_queryset_managers"][manager_fullname] = new_manager_info.fullname + # Add the new manager to the current module + module = api.modules[api.cur_mod_id] + module.names[name or manager_name] = SymbolTableNode( + GDEF, new_manager_info, plugin_generated=True, no_serialize=False + ) + return new_manager_info @@ -261,7 +261,7 @@ def create_manager_class( base_manager_instance = fill_typevars(base_manager_info) assert isinstance(base_manager_instance, Instance) - manager_info = api.basic_new_typeinfo(name, basetype_or_fallback=base_manager_instance, line=line) + manager_info = helpers.create_type_info(name, api.cur_mod_id, bases=[base_manager_instance]) manager_info.line = line manager_info.type_vars = base_manager_info.type_vars manager_info.defn.type_vars = base_manager_info.defn.type_vars diff --git a/mypy_django_plugin/transformers/models.py b/mypy_django_plugin/transformers/models.py index 23f138ecf1..befe90f2b2 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/mypy_django_plugin/transformers/models.py @@ -29,7 +29,6 @@ from mypy_django_plugin.errorcodes import MANAGER_MISSING from mypy_django_plugin.lib import fullnames, helpers from mypy_django_plugin.lib.fullnames import ANNOTATIONS_FULLNAME, ANY_ATTR_ALLOWED_CLASS_FULLNAME, MODEL_CLASS_FULLNAME -from mypy_django_plugin.lib.helpers import add_new_class_for_module from mypy_django_plugin.transformers import fields from mypy_django_plugin.transformers.fields import get_field_descriptor_types from mypy_django_plugin.transformers.managers import create_manager_info_from_from_queryset_call @@ -357,18 +356,7 @@ class MyModel(models.Model): if not isinstance(expr, CallExpr) or not isinstance(expr.callee, CallExpr): return None - new_manager_info = create_manager_info_from_from_queryset_call(self.api, expr.callee) - - if new_manager_info: - assert self.api.add_symbol_table_node( - # We'll use `new_manager_info.name` instead of `manager_class_name` here - # to handle possible name collisions, as it's unique. - new_manager_info.name, - # Note that the generated manager type is always inserted at module level - SymbolTableNode(MDEF, new_manager_info, plugin_generated=True), - ) - - return new_manager_info + return create_manager_info_from_from_queryset_call(self.api, expr.callee) class AddDefaultManagerAttribute(ModelClassInitializer): @@ -646,7 +634,7 @@ def get_or_create_annotated_type( else: annotated_model_type = api.named_generic_type(ANY_ATTR_ALLOWED_CLASS_FULLNAME, []) - annotated_typeinfo = add_new_class_for_module( + annotated_typeinfo = helpers.add_new_class_for_module( model_module_file, type_name, bases=[model_type] if fields_dict is not None else [model_type, annotated_model_type], diff --git a/tests/typecheck/managers/querysets/test_from_queryset.yml b/tests/typecheck/managers/querysets/test_from_queryset.yml index 467b11e298..5cdd47f98e 100644 --- a/tests/typecheck/managers/querysets/test_from_queryset.yml +++ b/tests/typecheck/managers/querysets/test_from_queryset.yml @@ -362,12 +362,14 @@ - case: test_queryset_in_model_class_body main: | from myapp.models import MyModel - reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.MyModel.MyManagerFromMyQuerySet[myapp.models.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.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.custom().filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.MyQuerySet" - reveal_type(MyModel.objects2) # N: Revealed type is "myapp.models.MyModel.MyManagerFromMyQuerySet[myapp.models.MyModel]" + 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]" installed_apps: - myapp files: @@ -387,6 +389,31 @@ objects = MyManager.from_queryset(MyQuerySet)() objects2 = MyManager.from_queryset(MyQuerySet)() +- case: test_queryset_in_model_class_body_subclass + main: | + from myapp.models import MyModel + reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.BaseManagerFromBaseQuerySet[myapp.models.MyModel]" + installed_apps: + - myapp + files: + - path: myapp/__init__.py + - path: myapp/models.py + content: | + from django.db import models + + class BaseManager(models.Manager["BaseModel"]): + pass + + class BaseQuerySet(models.QuerySet["BaseModel"]): + def custom(self) -> "BaseQuerySet": + pass + + class BaseModel(models.Model): + objects = BaseManager.from_queryset(BaseQuerySet)() + + class MyModel(BaseModel): + pass + - case: from_queryset_includes_methods_returning_queryset main: | from myapp.models import MyModel