fill QuerySet generics using the manager's model type

[asottile]

new test was previously failing with:

__________________ test_from_queryset_with_concrete_subclass ___________________
/tmp/django-stubs/tests/typecheck/managers/querysets/test_from_queryset.yml:895: 
E   pytest_mypy_plugins.utils.TypecheckAssertionError: Invalid output: 
E   Actual:
E     ...
E     main:4: note: Revealed type is "myapp.models.CustomQuerySet" (diff)
E     main:5: note: Revealed type is "D`2 = M`1"    (diff)
E   Expected:
E     ...
E     main:4: note: Revealed type is "myapp.models.CustomQuerySet[myapp.models.Concrete, myapp.models.Concrete]" (diff)
E     main:5: note: Revealed type is "myapp.models.Concrete" (diff)
E   Alignment of first line difference:
E     E: ... "myapp.models.CustomQuerySet[myapp.models.Concrete, myapp.models.Con...
E     A: ... "myapp.models.CustomQuerySet"...
E                                        ^

[asottile]

these (and all the similar ones) are more correct now -- these classes aren't generic so they shouldn't have type variables

I'm actually surprised mypy's internals don't enforce that an Instance has the correct number of typevars filled in!

[asottile]

nice, this was previously:

E     myapp/models:35: error: Need type annotation for "objects"  [var-annotated] (diff)

[asottile]

actually I have a better idea I want to try out for the as_manager stuff -- so hold off on merging for a bit :)

[flaeppe]

Nice work! I was fiddling around with those Var insertions in #2280 but it would've taken me ages to figure out to swap to a method hook.

actually I have a better idea I want to try out for the as_manager stuff -- so hold off on merging for a bit :)

Just let me know when you want me to have a look again

[asottile]

tried instead overriding the base class hook to add the as_manager method directly -- ended up being a dead end -- got close though so I might try the other approach again at some point:

diff --git a/mypy_django_plugin/main.py b/mypy_django_plugin/main.py
index 55455037..e9197324 100644
--- a/mypy_django_plugin/main.py
+++ b/mypy_django_plugin/main.py
@@ -37,8 +37,7 @@ from mypy_django_plugin.transformers import (
 )
 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,
+    add_as_manager_to_queryset_class,
     create_new_manager_class_from_from_queryset_method,
     reparametrize_any_manager_hook,
     resolve_manager_method,
@@ -209,10 +208,6 @@ class NewSemanalDjangoPlugin(Plugin):
                 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)
@@ -250,6 +245,11 @@ class NewSemanalDjangoPlugin(Plugin):
         # Base class is a Form class definition
         if fullname in self._get_current_form_bases():
             return transform_form_class
+
+        # Base class is a QuerySet class definition
+        sym = self.lookup_fully_qualified(fullname)
+        if sym is not None and isinstance(sym.node, TypeInfo) and sym.node.has_base(fullnames.QUERYSET_CLASS_FULLNAME):
+            return add_as_manager_to_queryset_class
         return None
 
     def get_attribute_hook(self, fullname: str) -> Optional[Callable[[AttributeContext], MypyType]]:
@@ -308,10 +308,6 @@ class NewSemanalDjangoPlugin(Plugin):
             info = self._get_typeinfo_or_none(class_name)
             if info and info.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME):
                 return create_new_manager_class_from_from_queryset_method
-        elif method_name == "as_manager":
-            info = self._get_typeinfo_or_none(class_name)
-            if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME):
-                return create_new_manager_class_from_as_manager_method
         return None
 
     def report_config_data(self, ctx: ReportConfigContext) -> Dict[str, Any]:
diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py
index ec13bf20..9404b5f8 100644
--- a/mypy_django_plugin/transformers/managers.py
+++ b/mypy_django_plugin/transformers/managers.py
@@ -16,7 +16,8 @@ from mypy.nodes import (
     SymbolTableNode,
     TypeInfo,
 )
-from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext, MethodContext
+from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext
+from mypy.plugins.common import add_method_to_class
 from mypy.semanal import SemanticAnalyzer
 from mypy.semanal_shared import has_placeholder
 from mypy.subtypes import find_member
@@ -407,15 +408,20 @@ def create_manager_info_from_from_queryset_call(
     return new_manager_info
 
 
-def create_manager_class(
-    api: SemanticAnalyzer, base_manager_info: TypeInfo, name: str, line: int, with_unique_name: bool
-) -> TypeInfo:
+def _base_manager_instance(base_manager_info: TypeInfo) -> Instance:
     base_manager_instance = fill_typevars(base_manager_info)
     assert isinstance(base_manager_instance, Instance)
 
     # If any of the type vars are undefined we need to defer. This is handled by the caller
     if any(has_placeholder(type_var) for type_var in base_manager_info.defn.type_vars):
         raise helpers.IncompleteDefnException
+    return base_manager_instance
+
+
+def create_manager_class(
+    api: SemanticAnalyzer, base_manager_info: TypeInfo, name: str, line: int, with_unique_name: bool
+) -> TypeInfo:
+    base_manager_instance = _base_manager_instance(base_manager_info)
 
     if with_unique_name:
         manager_info = helpers.add_new_class_for_module(
@@ -482,43 +488,26 @@ def populate_manager_from_queryset(manager_info: TypeInfo, queryset_info: TypeIn
         )
 
 
-def create_new_manager_class_from_as_manager_method(ctx: DynamicClassDefContext) -> None:
-    """
-    Insert a new manager class node for a
-
-    ```
-    <manager name> = <QuerySet>.as_manager()
-    ```
-    """
+def add_as_manager_to_queryset_class(ctx: ClassDefContext) -> None:
     semanal_api = helpers.get_semanal_api(ctx)
-    # 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 manager_node.type is not None:
-        # This is just a deferral run where our work is already finished
-        return
 
-    manager_sym = semanal_api.lookup_fully_qualified_or_none(fullnames.MANAGER_CLASS_FULLNAME)
-    assert manager_sym is not None
-    manager_base = manager_sym.node
-    if manager_base is None:
+    def _defer() -> None:
         if not semanal_api.final_iteration:
             semanal_api.defer()
-        return
-
-    assert isinstance(manager_base, TypeInfo)
 
-    callee = ctx.call.callee
-    assert isinstance(callee, MemberExpr)
-    assert isinstance(callee.expr, RefExpr)
-
-    queryset_info = callee.expr.node
+    queryset_info = semanal_api.type
     if queryset_info is None:
-        if not semanal_api.final_iteration:
-            semanal_api.defer()
+        return _defer()
+
+    # either a manual `as_manager` definition or this is a deferral pass
+    if "as_manager" in queryset_info.names:
         return
 
-    assert isinstance(queryset_info, TypeInfo)
+    manager_sym = semanal_api.lookup_fully_qualified_or_none(fullnames.MANAGER_CLASS_FULLNAME)
+    if manager_sym is None or not isinstance(manager_sym.node, TypeInfo):
+        return _defer()
 
+    manager_base = manager_sym.node
     manager_class_name = manager_base.name + "From" + queryset_info.name
     current_module = semanal_api.modules[semanal_api.cur_mod_id]
     existing_sym = current_module.names.get(manager_class_name)
@@ -537,7 +526,7 @@ def create_new_manager_class_from_as_manager_method(ctx: DynamicClassDefContext)
                 api=semanal_api,
                 base_manager_info=manager_base,
                 name=manager_class_name,
-                line=ctx.call.line,
+                line=queryset_info.line,
                 with_unique_name=True,
             )
         except helpers.IncompleteDefnException:
@@ -555,34 +544,22 @@ def create_new_manager_class_from_as_manager_method(ctx: DynamicClassDefContext)
         queryset_info.metadata["django_as_manager_names"][semanal_api.cur_mod_id] = new_manager_info.name
 
     # 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
-        # to handle possible name collisions, as it's unique.
-        new_manager_info.name,
+    # We'll use `new_manager_info.name` instead of `manager_class_name` here
+    # to handle possible name collisions, as it's unique.
+    added = semanal_api.modules[semanal_api.cur_mod_id].names[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(GDEF, new_manager_info, plugin_generated=True)
     )
     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)])
+    add_method_to_class(
+        semanal_api,
+        ctx.cls,
+        "as_manager",
+        args=[],
+        return_type=Instance(new_manager_info, [AnyType(TypeOfAny.from_omitted_generics)]),
+        is_classmethod=True,
+    )
 
 
 def reparametrize_any_manager_hook(ctx: ClassDefContext) -> None:

ended up hitting:

$ ../django-stubs/venv/bin/pip uninstall django-stubs -qqy && ../django-stubs/venv/bin/pip install -qq --no-deps ../django-stubs/ && ../django-stubs/venv/bin/python -m mypy main.py --disable-error-code empty-body --show-traceback
myapp/models.py:10: error: Return type "ManagerFromMyQuerySet[Any]" of "as_manager" incompatible with return type "ManagerFromBaseQuerySet[Any]" in supertype "BaseQuerySet"  [override]
main.py:2: note: Revealed type is "myapp.models.ManagerFromMyQuerySet[myapp.models.MyModel]"
Found 1 error in 1 file (checked 1 source file)

so yeah let's review / merge this in its current state for now!