From 5444f4518737b8a8bc01316c849da636f41ae583 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 | 30 ++---------------- .../managers/querysets/test_from_queryset.yml | 31 +++++++++++++++++-- 4 files changed, 57 insertions(+), 46 deletions(-) diff --git a/mypy_django_plugin/lib/helpers.py b/mypy_django_plugin/lib/helpers.py index da19db01f..abbdccb2d 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 6431524f2..ce41b6ee9 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 23f138ecf..a4d2f7a1d 100644 --- a/mypy_django_plugin/transformers/models.py +++ b/mypy_django_plugin/transformers/models.py @@ -5,19 +5,7 @@ from django.db.models.fields.related import ForeignKey from django.db.models.fields.reverse_related import ManyToManyRel, ManyToOneRel, OneToOneRel from mypy.checker import TypeChecker -from mypy.nodes import ( - ARG_STAR2, - MDEF, - Argument, - AssignmentStmt, - CallExpr, - Context, - FuncDef, - NameExpr, - SymbolTableNode, - TypeInfo, - Var, -) +from mypy.nodes import ARG_STAR2, Argument, AssignmentStmt, CallExpr, Context, FuncDef, NameExpr, TypeInfo, Var from mypy.plugin import AnalyzeTypeContext, AttributeContext, CheckerPluginInterface, ClassDefContext from mypy.plugins import common from mypy.semanal import SemanticAnalyzer @@ -29,7 +17,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 +344,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 +622,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 467b11e29..5cdd47f98 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