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.

Author: ljodal

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
     )

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

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],

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