construct as_manager instance at checker time

Author: asottile

mypy_django_plugin/main.py

@@ -37,6 +37,7 @@
 )
 from mypy_django_plugin.transformers.functional import resolve_str_promise_attribute
 from mypy_django_plugin.transformers.managers import (
+    construct_as_manager_instance,
     create_new_manager_class_from_as_manager_method,
     create_new_manager_class_from_from_queryset_method,
     reparametrize_any_manager_hook,
@@ -208,6 +209,10 @@ def get_method_hook(self, fullname: str) -> Optional[Callable[[MethodContext], M
                 fullnames.REVERSE_MANY_TO_ONE_DESCRIPTOR: manytoone.refine_many_to_one_related_manager,
             }
             return hooks.get(class_fullname)
+        elif method_name == "as_manager":
+            info = self._get_typeinfo_or_none(class_fullname)
+            if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME):
+                return partial(construct_as_manager_instance, info=info)
 
         if method_name in self.manager_and_queryset_method_hooks:
             info = self._get_typeinfo_or_none(class_fullname)

mypy_django_plugin/transformers/managers.py

@@ -15,9 +15,8 @@
     StrExpr,
     SymbolTableNode,
     TypeInfo,
-    Var,
 )
-from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext
+from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext, MethodContext
 from mypy.semanal import SemanticAnalyzer
 from mypy.semanal_shared import has_placeholder
 from mypy.subtypes import find_member
@@ -552,23 +551,9 @@ def create_new_manager_class_from_as_manager_method(ctx: DynamicClassDefContext)
             manager_name=manager_class_name,
             manager_base=manager_base,
         )
+        queryset_info.metadata.setdefault("django_as_manager_names", {})
+        queryset_info.metadata["django_as_manager_names"][semanal_api.cur_mod_id] = new_manager_info.name
 
-    # Whenever `<QuerySet>.as_manager()` isn't called at class level, we want to ensure
-    # that the variable is an instance of our generated manager. Instead of the return
-    # value of `.as_manager()`. Though model argument is populated as `Any`.
-    # `transformers.models.AddManagers` will populate a model's manager(s), when it
-    # finds it on class level.
-    var = Var(name=ctx.name, type=Instance(new_manager_info, [AnyType(TypeOfAny.from_omitted_generics)]))
-    var.info = new_manager_info
-    var._fullname = f"{current_module.fullname}.{ctx.name}"
-    var.is_inferred = True
-    # Note: Order of `add_symbol_table_node` calls matters. Depending on what level
-    # we've found the `.as_manager()` call. Point here being that we want to replace the
-    # `.as_manager` return value with our newly created manager.
-    added = semanal_api.add_symbol_table_node(
-        ctx.name, SymbolTableNode(semanal_api.current_symbol_kind(), var, plugin_generated=True)
-    )
-    assert added
     # Add the new manager to the current module
     added = semanal_api.add_symbol_table_node(
         # We'll use `new_manager_info.name` instead of `manager_class_name` here
@@ -580,6 +565,26 @@ def create_new_manager_class_from_as_manager_method(ctx: DynamicClassDefContext)
     assert added
 
 
+def construct_as_manager_instance(ctx: MethodContext, *, info: TypeInfo) -> MypyType:
+    api = helpers.get_typechecker_api(ctx)
+    module = helpers.get_current_module(api)
+    try:
+        manager_name = info.metadata["django_as_manager_names"][module.fullname]
+    except KeyError:
+        return ctx.default_return_type
+
+    manager_node = api.lookup(manager_name)
+    if not isinstance(manager_node.node, TypeInfo):
+        return ctx.default_return_type
+
+    # Whenever `<QuerySet>.as_manager()` isn't called at class level, we want to ensure
+    # that the variable is an instance of our generated manager. Instead of the return
+    # value of `.as_manager()`. Though model argument is populated as `Any`.
+    # `transformers.models.AddManagers` will populate a model's manager(s), when it
+    # finds it on class level.
+    return Instance(manager_node.node, [AnyType(TypeOfAny.from_omitted_generics)])
+
+
 def reparametrize_any_manager_hook(ctx: ClassDefContext) -> None:
     """
     Add implicit generics to manager classes that are defined without generic.

tests/typecheck/managers/querysets/test_as_manager.yml

@@ -28,7 +28,7 @@
                     def just_int(self) -> int: ...
 
                 class MyModel(models.Model):
-                    objects = MyQuerySet.as_manager()  # type: ignore[var-annotated]
+                    objects = MyQuerySet.as_manager()
 
                 class QuerySetWithoutSelf(models.QuerySet["MyModelWithoutSelf"]):
                     def method(self) -> "QuerySetWithoutSelf":
@@ -75,7 +75,7 @@
                     ...
 
                 class MyModel(models.Model):
-                    objects = MyQuerySet.as_manager()  # type: ignore[var-annotated]
+                    objects = MyQuerySet.as_manager()
 
 -   case: model_gets_generated_manager_as_default_manager
     main: |
@@ -285,7 +285,7 @@
                   objects = MyModelQuerySet.as_manager()
 
               class MyOtherModel(models.Model):
-                  objects = _MyModelQuerySet2.as_manager()  # type: ignore
+                  objects = _MyModelQuerySet2.as_manager()
 
 -   case: handles_type_vars
     main: |