Add test and other minor tweaks

mypy_django_plugin/transformers/managers.py

@@ -1,6 +1,5 @@
-from typing import Optional, Tuple, Union
+from typing import Optional, Union
 
-from django.db.models import base, manager
 from mypy.checker import TypeChecker, fill_typevars
 from mypy.nodes import (
     GDEF,
@@ -9,15 +8,14 @@
     FuncBase,
     FuncDef,
     MemberExpr,
-    NameExpr,
     OverloadedFuncDef,
     RefExpr,
     StrExpr,
     SymbolTableNode,
     TypeInfo,
     Var,
 )
-from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext, SemanticAnalyzerPluginInterface
+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
@@ -192,19 +190,12 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
         # This is just a deferral run where our work is already finished
         return
 
-    new_manager_info, generated_name = create_manager_info_from_from_queryset_call(ctx.api, ctx.call, ctx.name)
+    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 generated_name
-    manager_fullname = ".".join(["django.db.models.manager", generated_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
-
     # 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)
 
@@ -217,7 +208,7 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
 
 def create_manager_info_from_from_queryset_call(
     api: SemanticAnalyzerPluginInterface, call_expr: CallExpr, name: Optional[str] = None
-) -> Tuple[Optional[TypeInfo], Optional[str]]:
+) -> Optional[TypeInfo]:
     """
     Extract manager and queryset TypeInfo from a from_queryset call.
     """
@@ -236,14 +227,14 @@ def create_manager_info_from_from_queryset_call(
         or not isinstance(call_expr.args[0].node, TypeInfo)
         or not call_expr.args[0].node.has_base(fullnames.QUERYSET_CLASS_FULLNAME)
     ):
-        return None, None
+        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, None
+        return None
 
     if len(call_expr.args) == 2 and isinstance(call_expr.args[1], StrExpr):
         manager_name = call_expr.args[1].value
@@ -254,7 +245,13 @@ def create_manager_info_from_from_queryset_call(
 
     popuplate_manager_from_queryset(new_manager_info, queryset_info)
 
-    return new_manager_info, manager_name
+    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
+
+    return new_manager_info
 
 
 def create_manager_class(

mypy_django_plugin/transformers/models.py

@@ -7,7 +7,7 @@
 from mypy.checker import TypeChecker
 from mypy.nodes import (
     ARG_STAR2,
-    GDEF,
+    MDEF,
     Argument,
     AssignmentStmt,
     CallExpr,
@@ -24,7 +24,6 @@
 from mypy.types import AnyType, Instance
 from mypy.types import Type as MypyType
 from mypy.types import TypedDictType, TypeOfAny
-from typing_extensions import reveal_type
 
 from mypy_django_plugin.django.context import DjangoContext
 from mypy_django_plugin.errorcodes import MANAGER_MISSING
@@ -317,18 +316,15 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
 
     def get_manager_expression(self, name: str) -> Optional[AssignmentStmt]:
         # TODO: What happens if the manager is defined multiple times?
-        return next(
-            (
-                expr
-                for expr in self.ctx.cls.defs.body
-                if (
-                    isinstance(expr, AssignmentStmt)
-                    and isinstance(expr.lvalues[0], NameExpr)
-                    and expr.lvalues[0].name == name
-                )
-            ),
-            None,
-        )
+        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]:
         """
@@ -361,15 +357,15 @@ 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)
+        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(GDEF, new_manager_info, plugin_generated=True),
+                SymbolTableNode(MDEF, new_manager_info, plugin_generated=True),
             )
 
         return new_manager_info
@@ -382,6 +378,7 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
 
         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:

tests/typecheck/managers/querysets/test_from_queryset.yml

@@ -359,42 +359,33 @@
                 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.MyModel.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]"
     installed_apps:
         - myapp
     files:
         -   path: myapp/__init__.py
         -   path: myapp/models.py
             content: |
                 from django.db import models
-                from django.db.models.manager import BaseManager
 
-                class MyQuerySet(models.QuerySet["MyModel"]):
-                    def queryset_method(self) -> int:
-                        return 1
+                class MyManager(models.Manager["MyModel"]):
+                    pass
 
-                class MyManager(BaseManager):
-                    ...
+                class MyQuerySet(models.QuerySet["MyModel"]):
+                    def custom(self) -> "MyQuerySet":
+                        pass
 
-                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: from_queryset_includes_methods_returning_queryset
     main: |