Add support for inline from_queryset in model classes (#1045)

* Add support for inline from_queryset in model classes

This adds support for calling <Manager>.from_queryset(<QuerySet>)()
inline in models, for example like this:

    class MyModel(models.Model):
        objects = MyManager.from_queryset(MyQuerySet)()

This is done by inspecting the class body in the transform_class_hook

* Fix missing methods on copied manager

* Add test and other minor tweaks

* 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

@@ -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
     )
@@ -382,29 +390,25 @@ def copy_method_to_another_class(
     method_node: FuncDef,
     return_type: Optional[MypyType] = None,
     original_module_name: Optional[str] = None,
-) -> None:
+) -> bool:
     semanal_api = get_semanal_api(ctx)
     if method_node.type is None:
-        if not semanal_api.final_iteration:
-            semanal_api.defer()
-            return
-
         arguments, return_type = build_unannotated_method_args(method_node)
         add_method_to_class(
             semanal_api, ctx.cls, new_method_name, args=arguments, return_type=return_type, self_type=self_type
         )
-        return
+        return True
 
     method_type = method_node.type
     if not isinstance(method_type, CallableType):
         if not semanal_api.final_iteration:
             semanal_api.defer()
-        return
+        return False
 
     if return_type is None:
         return_type = bind_or_analyze_type(method_type.ret_type, semanal_api, original_module_name)
     if return_type is None:
-        return
+        return False
 
     # We build the arguments from the method signature (`CallableType`), because if we were to
     # use the arguments from the method node (`FuncDef.arguments`) we're not compatible with
@@ -417,7 +421,7 @@ def copy_method_to_another_class(
     ):
         bound_arg_type = bind_or_analyze_type(arg_type, semanal_api, original_module_name)
         if bound_arg_type is None:
-            return
+            return False
         if arg_name is None and hasattr(method_node, "arguments"):
             arg_name = method_node.arguments[pos].variable.name
         arguments.append(
@@ -435,6 +439,8 @@ def copy_method_to_another_class(
         semanal_api, ctx.cls, new_method_name, args=arguments, return_type=return_type, self_type=self_type
     )
 
+    return True
+
 
 def add_new_manager_base(api: SemanticAnalyzerPluginInterface, fullname: str) -> None:
     sym = api.lookup_fully_qualified_or_none(fullnames.MANAGER_CLASS_FULLNAME)

mypy_django_plugin/main.py

@@ -24,7 +24,6 @@ from mypy_django_plugin.lib import fullnames, helpers
 from mypy_django_plugin.transformers import fields, forms, init_create, meta, querysets, request, settings
 from mypy_django_plugin.transformers.managers import (
     create_new_manager_class_from_from_queryset_method,
-    fail_if_manager_type_created_in_model_body,
     resolve_manager_method,
 )
 from mypy_django_plugin.transformers.models import (
@@ -237,11 +236,6 @@ class NewSemanalDjangoPlugin(Plugin):
                 django_context=self.django_context,
             )
 
-        elif method_name == "from_queryset":
-            info = self._get_typeinfo_or_none(class_fullname)
-            if info and info.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME):
-                return fail_if_manager_type_created_in_model_body
-
         return None
 
     def get_base_class_hook(self, fullname: str) -> Optional[Callable[[ClassDefContext], None]]:

mypy_django_plugin/transformers/managers.py

@@ -8,7 +8,6 @@ from mypy.nodes import (
     FuncBase,
     FuncDef,
     MemberExpr,
-    NameExpr,
     OverloadedFuncDef,
     RefExpr,
     StrExpr,
@@ -16,13 +15,12 @@ from mypy.nodes import (
     TypeInfo,
     Var,
 )
-from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext, MethodContext
+from mypy.plugin import AttributeContext, DynamicClassDefContext, SemanticAnalyzerPluginInterface
 from mypy.types import AnyType, CallableType, Instance, ProperType
 from mypy.types import Type as MypyType
 from mypy.types import TypeOfAny
 from typing_extensions import Final
 
-from mypy_django_plugin import errorcodes
 from mypy_django_plugin.lib import fullnames, helpers
 
 MANAGER_METHODS_RETURNING_QUERYSET: Final = frozenset(
@@ -182,81 +180,110 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
     """
     semanal_api = helpers.get_semanal_api(ctx)
 
+    # TODO: Emit an error when called in a class scope
+    if semanal_api.is_class_scope():
+        return
+
     # Don't redeclare the manager class if we've already defined it.
     manager_node = semanal_api.lookup_current_scope(ctx.name)
     if manager_node and isinstance(manager_node.node, TypeInfo):
         # This is just a deferral run where our work is already finished
         return
 
-    callee = ctx.call.callee
-    assert isinstance(callee, MemberExpr)
-    assert isinstance(callee.expr, RefExpr)
-
-    base_manager_info = callee.expr.node
-    if base_manager_info is None:
-        if not semanal_api.final_iteration:
-            semanal_api.defer()
+    new_manager_info = create_manager_info_from_from_queryset_call(ctx.api, ctx.call, ctx.name)
+    if new_manager_info is None:
+        if not ctx.api.final_iteration:
+            ctx.api.defer()
         return
 
-    assert isinstance(base_manager_info, TypeInfo)
-
-    passed_queryset = ctx.call.args[0]
-    assert isinstance(passed_queryset, NameExpr)
-
-    derived_queryset_fullname = passed_queryset.fullname
-    if derived_queryset_fullname is None:
-        # In some cases, due to the way the semantic analyzer works, only passed_queryset.name is available.
-        # But it should be analyzed again, so this isn't a problem.
-        return
-
-    base_manager_instance = fill_typevars(base_manager_info)
-    assert isinstance(base_manager_instance, Instance)
-    new_manager_info = semanal_api.basic_new_typeinfo(
-        ctx.name, basetype_or_fallback=base_manager_instance, line=ctx.call.line
-    )
-
-    sym = semanal_api.lookup_fully_qualified_or_none(derived_queryset_fullname)
-    assert sym is not None
-    if sym.node is None:
-        if not semanal_api.final_iteration:
-            semanal_api.defer()
-        else:
-            # inherit from Any to prevent false-positives, if queryset class cannot be resolved
-            new_manager_info.fallback_to_any = True
-        return
-
-    derived_queryset_info = sym.node
-    assert isinstance(derived_queryset_info, TypeInfo)
-
-    new_manager_info.line = ctx.call.line
-    new_manager_info.type_vars = base_manager_info.type_vars
-    new_manager_info.defn.type_vars = base_manager_info.defn.type_vars
-    new_manager_info.defn.line = ctx.call.line
-    new_manager_info.metaclass_type = new_manager_info.calculate_metaclass_type()
-    # Stash the queryset fullname which was passed to .from_queryset
-    # So that our 'resolve_manager_method' attribute hook can fetch the method from that QuerySet class
-    new_manager_info.metadata["django"] = {"from_queryset_manager": derived_queryset_fullname}
-
-    if len(ctx.call.args) > 1:
-        expr = ctx.call.args[1]
-        assert isinstance(expr, StrExpr)
-        custom_manager_generated_name = expr.value
-    else:
-        custom_manager_generated_name = base_manager_info.name + "From" + derived_queryset_info.name
-
-    custom_manager_generated_fullname = ".".join(["django.db.models.manager", custom_manager_generated_name])
-    base_manager_info.metadata.setdefault("from_queryset_managers", {})
-    base_manager_info.metadata["from_queryset_managers"][custom_manager_generated_fullname] = new_manager_info.fullname
-
     # 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)
 
-    class_def_context = ClassDefContext(cls=new_manager_info.defn, reason=ctx.call, api=semanal_api)
-    self_type = fill_typevars(new_manager_info)
-    assert isinstance(self_type, Instance)
+
+def create_manager_info_from_from_queryset_call(
+    api: SemanticAnalyzerPluginInterface, call_expr: CallExpr, name: Optional[str] = None
+) -> Optional[TypeInfo]:
+    """
+    Extract manager and queryset TypeInfo from a from_queryset call.
+    """
+
+    if (
+        # Check that this is a from_queryset call on a manager subclass
+        not isinstance(call_expr.callee, MemberExpr)
+        or not isinstance(call_expr.callee.expr, RefExpr)
+        or not isinstance(call_expr.callee.expr.node, TypeInfo)
+        or not call_expr.callee.expr.node.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME)
+        or not call_expr.callee.name == "from_queryset"
+        # Check that the call has one or two arguments and that the first is a
+        # QuerySet subclass
+        or not 1 <= len(call_expr.args) <= 2
+        or not isinstance(call_expr.args[0], RefExpr)
+        or not isinstance(call_expr.args[0].node, TypeInfo)
+        or not call_expr.args[0].node.has_base(fullnames.QUERYSET_CLASS_FULLNAME)
+    ):
+        return None
+
+    base_manager_info, queryset_info = call_expr.callee.expr.node, call_expr.args[0].node
+    if queryset_info.fullname is None:
+        # In some cases, due to the way the semantic analyzer works, only
+        # passed_queryset.name is available. But it should be analyzed again,
+        # so this isn't a problem.
+        return None
+
+    if len(call_expr.args) == 2 and isinstance(call_expr.args[1], StrExpr):
+        manager_name = call_expr.args[1].value
+    else:
+        manager_name = f"{base_manager_info.name}From{queryset_info.name}"
+
+    new_manager_info = create_manager_class(api, base_manager_info, name or manager_name, call_expr.line)
+
+    popuplate_manager_from_queryset(new_manager_info, queryset_info)
+
+    manager_fullname = ".".join(["django.db.models.manager", manager_name])
+
+    base_manager_info = new_manager_info.mro[1]
+    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
+
+
+def create_manager_class(
+    api: SemanticAnalyzerPluginInterface, base_manager_info: TypeInfo, name: str, line: int
+) -> TypeInfo:
+
+    base_manager_instance = fill_typevars(base_manager_info)
+    assert isinstance(base_manager_instance, Instance)
+
+    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
+    manager_info.defn.line = line
+    manager_info.metaclass_type = manager_info.calculate_metaclass_type()
+
+    return manager_info
+
+
+def popuplate_manager_from_queryset(manager_info: TypeInfo, queryset_info: TypeInfo) -> None:
+    """
+    Add methods from the QuerySet class to the manager.
+    """
+
+    # Stash the queryset fullname which was passed to .from_queryset So that
+    # our 'resolve_manager_method' attribute hook can fetch the method from
+    # that QuerySet class
+    django_metadata = helpers.get_django_metadata(manager_info)
+    django_metadata["from_queryset_manager"] = queryset_info.fullname
 
     # We collect and mark up all methods before django.db.models.query.QuerySet as class members
-    for class_mro_info in derived_queryset_info.mro:
+    for class_mro_info in queryset_info.mro:
         if class_mro_info.fullname == fullnames.QUERYSET_CLASS_FULLNAME:
             break
         for name, sym in class_mro_info.names.items():
@@ -270,39 +297,19 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
             #    queryset_method: Any = ...
             #
             helpers.add_new_sym_for_info(
-                new_manager_info,
+                manager_info,
                 name=name,
                 sym_type=AnyType(TypeOfAny.special_form),
             )
 
-    # For methods on BaseManager that return a queryset we need to update the
-    # return type to be the actual queryset subclass used. This is done by
-    # adding the methods as attributes with type Any to the manager class,
-    # similar to how custom queryset methods are handled above. The actual type
-    # of these methods are resolved in resolve_manager_method.
-    for name in MANAGER_METHODS_RETURNING_QUERYSET:
+    # For methods on BaseManager that return a queryset we need to update
+    # the return type to be the actual queryset subclass used. This is done
+    # by adding the methods as attributes with type Any to the manager
+    # class. The actual type of these methods are resolved in
+    # resolve_manager_method.
+    for method_name in MANAGER_METHODS_RETURNING_QUERYSET:
         helpers.add_new_sym_for_info(
-            new_manager_info,
-            name=name,
+            manager_info,
+            name=method_name,
             sym_type=AnyType(TypeOfAny.special_form),
         )
-
-    # Insert the new manager (dynamic) class
-    assert semanal_api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, new_manager_info, plugin_generated=True))
-
-
-def fail_if_manager_type_created_in_model_body(ctx: MethodContext) -> MypyType:
-    """
-    Method hook that checks if method `<Manager>.from_queryset` is called inside a model class body.
-
-    Doing so won't, for instance, trigger the dynamic class hook(`create_new_manager_class_from_from_queryset_method`)
-    for managers.
-    """
-    api = helpers.get_typechecker_api(ctx)
-    outer_model_info = api.scope.active_class()
-    if not outer_model_info or not outer_model_info.has_base(fullnames.MODEL_CLASS_FULLNAME):
-        # Not inside a model class definition
-        return ctx.default_return_type
-
-    api.fail("`.from_queryset` called from inside model class body", ctx.context, code=errorcodes.MANAGER_UNTYPED)
-    return ctx.default_return_type

mypy_django_plugin/transformers/models.py

@@ -1,11 +1,11 @@
 from typing import Dict, List, Optional, Type, Union, cast
 
-from django.db.models.base import Model
+from django.db.models import Manager, Model
 from django.db.models.fields import DateField, DateTimeField, Field
 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, Argument, AssignmentStmt, Context, FuncDef, NameExpr, 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
@@ -17,9 +17,9 @@ from mypy_django_plugin.django.context import DjangoContext
 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
 
 
 class ModelClassInitializer:
@@ -198,12 +198,12 @@ class AddManagers(ModelClassInitializer):
         bases = []
         for original_base in base_manager_info.bases:
             if self.is_any_parametrized_manager(original_base):
-                if original_base.type is None:
-                    raise helpers.IncompleteDefnException()
-
                 original_base = helpers.reparametrize_instance(original_base, [Instance(self.model_classdef.info, [])])
             bases.append(original_base)
 
+        # TODO: This adds the manager to the module, even if we end up
+        # deferring. That can be avoided by not adding it to the module first,
+        # but rather waiting until we know we won't defer
         new_manager_info = self.add_new_class_for_current_module(name, bases)
         # copy fields to a new manager
         new_cls_def_context = ClassDefContext(cls=new_manager_info.defn, reason=self.ctx.reason, api=self.api)
@@ -212,13 +212,15 @@ class AddManagers(ModelClassInitializer):
         for name, sym in base_manager_info.names.items():
             # replace self type with new class, if copying method
             if isinstance(sym.node, FuncDef):
-                helpers.copy_method_to_another_class(
+                copied_method = helpers.copy_method_to_another_class(
                     new_cls_def_context,
                     self_type=custom_manager_type,
                     new_method_name=name,
                     method_node=sym.node,
                     original_module_name=base_manager_info.module_name,
                 )
+                if not copied_method and not self.api.final_iteration:
+                    raise helpers.IncompleteDefnException()
                 continue
 
             new_sym = sym.copy()
@@ -236,23 +238,23 @@ class AddManagers(ModelClassInitializer):
 
         incomplete_manager_defs = set()
         for manager_name, manager in model_cls._meta.managers_map.items():
+            # If the manager is already typed do nothing
+            manager_node = self.model_classdef.info.names.get(manager_name, None)
+            if manager_node and manager_node.type is not None:
+                continue
+
             manager_class_name = manager.__class__.__name__
             manager_fullname = helpers.get_class_fullname(manager.__class__)
-            try:
-                manager_info = self.lookup_typeinfo_or_incomplete_defn_error(manager_fullname)
-            except helpers.IncompleteDefnException as exc:
-                # Check if manager is a generated (dynamic class) manager
-                base_manager_fullname = helpers.get_class_fullname(manager.__class__.__bases__[0])
-                generated_managers = self.get_generated_manager_mappings(base_manager_fullname)
-                if manager_fullname not in generated_managers:
-                    # Manager doesn't appear to be generated. Track that we encountered an
-                    # incomplete definition and skip
-                    incomplete_manager_defs.add(manager_name)
-                    continue
 
-                manager_info = self.lookup_typeinfo(generated_managers[manager_fullname])
-                if manager_info is None:
-                    continue
+            manager_info = self.lookup_typeinfo(manager_fullname)
+            if manager_info is None:
+                manager_info = self.create_manager_from_from_queryset(manager_name)
+            if manager_info is None:
+                manager_info = self.get_dynamic_manager(manager_fullname, manager)
+
+            if manager_info is None:
+                incomplete_manager_defs.add(manager_name)
+                continue
 
             is_dynamically_generated = manager_info.metadata.get("django", {}).get("from_queryset_manager") is not None
             if manager_name not in self.model_classdef.info.names or is_dynamically_generated:
@@ -264,18 +266,20 @@ class AddManagers(ModelClassInitializer):
                 # related manager.
                 custom_model_manager_name = manager.model.__name__ + "_" + manager_class_name
                 try:
-                    custom_manager_type = self.create_new_model_parametrized_manager(
+                    manager_type = self.create_new_model_parametrized_manager(
                         custom_model_manager_name, base_manager_info=manager_info
                     )
                 except helpers.IncompleteDefnException:
                     continue
 
-                self.add_new_node_to_model_class(manager_name, custom_manager_type)
+                self.add_new_node_to_model_class(manager_name, manager_type)
+
+        if incomplete_manager_defs:
+            if not self.api.final_iteration:
+                #  Unless we're on the final round, see if another round could
+                #  figure out all manager types
+                raise helpers.IncompleteDefnException()
 
-        if incomplete_manager_defs and not self.api.final_iteration:
-            #  Unless we're on the final round, see if another round could figure out all manager types
-            raise helpers.IncompleteDefnException()
-        elif self.api.final_iteration:
             for manager_name in incomplete_manager_defs:
                 # We act graceful and set the type as the bare minimum we know of
                 # (Django's default) before finishing. And emit an error, to allow for
@@ -289,22 +293,59 @@ class AddManagers(ModelClassInitializer):
                     manager_name, Instance(django_manager_info, [Instance(self.model_classdef.info, [])])
                 )
                 # Find expression for e.g. `objects = SomeManager()`
-                manager_expr = [
-                    expr
-                    for expr in self.ctx.cls.defs.body
-                    if (
-                        isinstance(expr, AssignmentStmt)
-                        and isinstance(expr.lvalues[0], NameExpr)
-                        and expr.lvalues[0].name == manager_name
-                    )
-                ]
+                manager_expr = self.get_manager_expression(manager_name)
                 manager_fullname = f"{self.model_classdef.fullname}.{manager_name}"
                 self.api.fail(
                     f'Could not resolve manager type for "{manager_fullname}"',
-                    manager_expr[0] if manager_expr else self.ctx.cls,
+                    manager_expr if manager_expr else self.ctx.cls,
                     code=MANAGER_MISSING,
                 )
 
+    def get_manager_expression(self, name: str) -> Optional[AssignmentStmt]:
+        # TODO: What happens if the manager is defined multiple times?
+        for expr in self.ctx.cls.defs.body:
+            if (
+                isinstance(expr, AssignmentStmt)
+                and isinstance(expr.lvalues[0], NameExpr)
+                and expr.lvalues[0].name == name
+            ):
+                return expr
+
+        return None
+
+    def get_dynamic_manager(self, fullname: str, manager: Manager) -> Optional[TypeInfo]:
+        """
+        Try to get a dynamically defined manager
+        """
+
+        # Check if manager is a generated (dynamic class) manager
+        base_manager_fullname = helpers.get_class_fullname(manager.__class__.__bases__[0])
+        generated_managers = self.get_generated_manager_mappings(base_manager_fullname)
+
+        generated_manager_name: Optional[str] = generated_managers.get(fullname, None)
+        if generated_manager_name is None:
+            return None
+
+        return self.lookup_typeinfo(generated_manager_name)
+
+    def create_manager_from_from_queryset(self, name: str) -> Optional[TypeInfo]:
+        """
+        Try to create a manager from a .from_queryset call:
+
+            class MyModel(models.Model):
+                objects = MyManager.from_queryset(MyQuerySet)()
+        """
+
+        assign_statement = self.get_manager_expression(name)
+        if assign_statement is None:
+            return None
+
+        expr = assign_statement.rvalue
+        if not isinstance(expr, CallExpr) or not isinstance(expr.callee, CallExpr):
+            return None
+
+        return create_manager_info_from_from_queryset_call(self.api, expr.callee)
+
 
 class AddDefaultManagerAttribute(ModelClassInitializer):
     def run_with_model_cls(self, model_cls: Type[Model]) -> None:
@@ -313,6 +354,7 @@ class AddDefaultManagerAttribute(ModelClassInitializer):
 
         default_manager_cls = model_cls._meta.default_manager.__class__
         default_manager_fullname = helpers.get_class_fullname(default_manager_cls)
+
         try:
             default_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(default_manager_fullname)
         except helpers.IncompleteDefnException as exc:
@@ -580,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/fields/test_related.yml

@@ -678,7 +678,6 @@
                     def custom(self) -> None:
                         pass
 
-                # Note, that we cannot resolve dynamic calls for custom managers:
                 class Transaction(models.Model):
                     objects = BaseManager.from_queryset(TransactionQuerySet)
                     def test(self) -> None:
@@ -689,9 +688,8 @@
                 class TransactionLog(models.Model):
                     transaction = models.ForeignKey(Transaction, on_delete=models.CASCADE)
     out: |
-        myapp/models:9: error: `.from_queryset` called from inside model class body
-        myapp/models:11: note: Revealed type is "django.db.models.manager.RelatedManager[myapp.models.TransactionLog]"
-        myapp/models:13: note: Revealed type is "Any"
+        myapp/models:10: note: Revealed type is "django.db.models.manager.RelatedManager[myapp.models.TransactionLog]"
+        myapp/models:12: note: Revealed type is "Any"
 
 
 -   case: resolve_primary_keys_for_foreign_keys_with_abstract_self_model

tests/typecheck/managers/querysets/test_from_queryset.yml

@@ -359,12 +359,17 @@
                 class MyModel(models.Model):
                     objects = NewManager()
 
--   case: from_queryset_in_model_class_body_yields_message
+-   case: test_queryset_in_model_class_body
     main: |
         from myapp.models import MyModel
-        reveal_type(MyModel.base_manager)  # N: Revealed type is "myapp.models.BaseManagerFromMyQuerySet[myapp.models.MyModel]"
-        reveal_type(MyModel.manager)  # N: Revealed type is "myapp.models.ManagerFromMyQuerySet[myapp.models.MyModel]"
-        reveal_type(MyModel.custom_manager)  # N: Revealed type is "myapp.models.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.MyManagerFromMyQuerySet[myapp.models.MyModel]"
+        reveal_type(MyModel._default_manager)  # N: Revealed type is "myapp.models.MyManagerFromMyQuerySet[myapp.models.MyModel]"
     installed_apps:
         - myapp
     files:
@@ -372,29 +377,42 @@
         -   path: myapp/models.py
             content: |
                 from django.db import models
-                from django.db.models.manager import BaseManager
+
+                class MyManager(models.Manager["MyModel"]):
+                    pass
 
                 class MyQuerySet(models.QuerySet["MyModel"]):
-                    def queryset_method(self) -> int:
-                        return 1
+                    def custom(self) -> "MyQuerySet":
+                        pass
 
-                class MyManager(BaseManager):
-                    ...
-
-                BaseManagerFromMyQuerySet = BaseManager.from_queryset(MyQuerySet)
-                ManagerFromMyQuerySet = models.Manager.from_queryset(MyQuerySet)
-                MyManagerFromMyQuerySet = MyManager.from_queryset(MyQuerySet)
                 class MyModel(models.Model):
-                    objects1 = BaseManager.from_queryset(MyQuerySet)()  # E: `.from_queryset` called from inside model class body
-                    objects2 = BaseManager.from_queryset(MyQuerySet)  # E: `.from_queryset` called from inside model class body
-                    objects3 = models.Manager.from_queryset(MyQuerySet)()  # E: `.from_queryset` called from inside model class body
-                    objects4 = models.Manager.from_queryset(MyQuerySet)  # E: `.from_queryset` called from inside model class body
-                    objects5 = MyManager.from_queryset(MyQuerySet)  # E: `.from_queryset` called from inside model class body
-                    objects6 = MyManager.from_queryset(MyQuerySet)()  # E: `.from_queryset` called from inside model class body
-                    # Initiating the manager type is fine
-                    base_manager = BaseManagerFromMyQuerySet()
-                    manager = ManagerFromMyQuerySet()
-                    custom_manager = MyManagerFromMyQuerySet()
+                    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: |

tests/typecheck/managers/test_managers.yml

@@ -332,14 +332,14 @@
 -   case: custom_manager_returns_proper_model_types
     main: |
         from myapp.models import User
-        reveal_type(User.objects)  # N: Revealed type is "myapp.models.User_MyManager2[myapp.models.User]"
+        reveal_type(User.objects)  # N: Revealed type is "myapp.models.User_MyManager[myapp.models.User]"
         reveal_type(User.objects.select_related())  # N: Revealed type is "django.db.models.query._QuerySet[myapp.models.User, myapp.models.User]"
         reveal_type(User.objects.get())  # N: Revealed type is "myapp.models.User"
         reveal_type(User.objects.get_instance())  # N: Revealed type is "builtins.int"
         reveal_type(User.objects.get_instance_untyped('hello'))  # N: Revealed type is "Any"
 
         from myapp.models import ChildUser
-        reveal_type(ChildUser.objects)  # N: Revealed type is "myapp.models.ChildUser_MyManager2[myapp.models.ChildUser]"
+        reveal_type(ChildUser.objects)  # N: Revealed type is "myapp.models.ChildUser_MyManager[myapp.models.ChildUser]"
         reveal_type(ChildUser.objects.select_related())  # N: Revealed type is "django.db.models.query._QuerySet[myapp.models.ChildUser, myapp.models.ChildUser]"
         reveal_type(ChildUser.objects.get())  # N: Revealed type is "myapp.models.ChildUser"
         reveal_type(ChildUser.objects.get_instance())  # N: Revealed type is "builtins.int"