Support inheriting manager classes created via .from_queryset

mypy_django_plugin/transformers/managers.py

@@ -3,23 +3,36 @@
 from mypy.checker import TypeChecker
 from mypy.nodes import (
     GDEF,
+    AssignmentStmt,
     CallExpr,
     Decorator,
     FuncBase,
     FuncDef,
     MemberExpr,
+    NameExpr,
     Node,
     OverloadedFuncDef,
     RefExpr,
     StrExpr,
     SymbolTableNode,
+    TypeAlias,
     TypeInfo,
     Var,
 )
 from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext
 from mypy.semanal import SemanticAnalyzer
 from mypy.semanal_shared import has_placeholder
-from mypy.types import AnyType, CallableType, FunctionLike, Instance, Overloaded, ProperType, TypeOfAny, TypeVarType
+from mypy.types import (
+    AnyType,
+    CallableType,
+    FunctionLike,
+    Instance,
+    Overloaded,
+    PlaceholderType,
+    ProperType,
+    TypeOfAny,
+    TypeVarType,
+)
 from mypy.types import Type as MypyType
 from mypy.typevars import fill_typevars
 
@@ -58,10 +71,11 @@ def get_method_type_from_dynamic_manager(
     Attempt to resolve a method on a manager that was built from '.from_queryset'
     """
 
-    manager_type_info = manager_instance.type
+    manager_type_info = manager_instance.type.get_containing_type_info(method_name)
 
     if (
-        "django" not in manager_type_info.metadata
+        manager_type_info is None
+        or "django" not in manager_type_info.metadata
         or "from_queryset_manager" not in manager_type_info.metadata["django"]
     ):
         # Manager isn't dynamically added
@@ -347,24 +361,73 @@ def create_manager_info_from_from_queryset_call(
     """
     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 call_expr.callee.name == "from_queryset"
         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"
+    ):
+        return None
+
+    module = api.modules[api.cur_mod_id]
+    base_manager_info = call_expr.callee.expr.node
+    # Mypy parses `NewManager = <Manager>.from_queryset(<QuerySet>)` as a `Var`
+    # assignment, but what is happening runtime is rather that `NewManager` becomes an
+    # alias to another type. But mypy can't see that type. Before anything else happens,
+    # we want to adjust this. But we don't have any type to alias yet and we can't be
+    # sure we _can_ have it yet either. Due to the state of eventual dependencies we
+    # might need to defer, so to begin with we insert a placeholder type.
+    if name is not None:  # New manager is aliased
+        assignment_symbol = module.names.get(name)
+        if (
+            assignment_symbol is not None
+            and assignment_symbol.fullname is not None
+            and not assignment_symbol.plugin_generated
+            and not isinstance(assignment_symbol.node, TypeAlias)
+        ):
+            assignment_stmt = next(
+                (
+                    stmt
+                    for stmt in module.defs
+                    if (
+                        isinstance(stmt, AssignmentStmt)
+                        and len(stmt.lvalues) == 1
+                        and isinstance(stmt.lvalues[0], NameExpr)
+                        and stmt.lvalues[0].name == name
+                    )
+                ),
+                None,
+            )
+            if assignment_stmt is not None:
+                # Mypy don't allow call expressions to be aliases, but to get any kind
+                # of handle to the invisible type we spoof mypy that this statement is a
+                # type alias..
+                assignment_stmt.is_alias_def = True
+                alias = TypeAlias(
+                    target=PlaceholderType(
+                        fullname=assignment_symbol.fullname,
+                        args=list(base_manager_info.defn.type_vars),
+                        line=call_expr.line,
+                    ),
+                    fullname=assignment_symbol.fullname,
+                    line=call_expr.line,
+                    column=0,
+                )
+                module.names[name] = SymbolTableNode(GDEF, alias, plugin_generated=True)
+
+    if (
         # 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
+        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
+    queryset_info = 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,
@@ -408,10 +471,19 @@ def create_manager_info_from_from_queryset_call(
 
     # Add the new manager to the current module
     # TODO: use proper SemanticAnalyzer API for that.
-    module = api.modules[api.cur_mod_id]
-    if name is not None and name != new_manager_info.name:
-        # Unless names are equal, there's 2 symbol names that needs the manager info
-        module.names[name] = SymbolTableNode(GDEF, new_manager_info, plugin_generated=True)
+    if name is not None:
+        assignment_symbol = module.names.get(name)
+        if (
+            assignment_symbol is not None
+            and name != new_manager_info.name
+            and assignment_symbol.plugin_generated
+            and isinstance(assignment_symbol.node, TypeAlias)
+            and isinstance(assignment_symbol.node.target, PlaceholderType)
+        ):
+            # Unless names are equal, there's 2 symbol names that needs to be updated
+            manager_instance = fill_typevars(new_manager_info)
+            assert isinstance(manager_instance, Instance)
+            assignment_symbol.node.target = manager_instance
 
     module.names[new_manager_info.name] = SymbolTableNode(GDEF, new_manager_info, plugin_generated=True)
     return new_manager_info

tests/typecheck/fields/test_related.yml

@@ -642,8 +642,8 @@
         reveal_type(user.article_set) # N: Revealed type is "myapp.models.Article_RelatedManager"
         reveal_type(user.book_set.add)  # N: Revealed type is "def (*objs: Union[myapp.models.Book, builtins.int], *, bulk: builtins.bool =)"
         reveal_type(user.article_set.add)  # N: Revealed type is "def (*objs: Union[myapp.models.Article, builtins.int], *, bulk: builtins.bool =)"
-        reveal_type(user.book_set.filter)  # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.LibraryEntityQuerySet[myapp.models.Book]"
-        reveal_type(user.article_set.filter)  # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.LibraryEntityQuerySet[myapp.models.Article]"
+        reveal_type(user.book_set.filter)  # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.LibraryEntityQuerySet"
+        reveal_type(user.article_set.filter)  # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.LibraryEntityQuerySet"
         reveal_type(user.book_set.queryset_method())  # N: Revealed type is "builtins.int"
         reveal_type(user.article_set.queryset_method())  # N: Revealed type is "builtins.int"
     installed_apps:
@@ -798,11 +798,11 @@
         from myapp.models.user import User
         reveal_type(Store().purchases)  # N: Revealed type is "myapp.models.purchase.Purchase_RelatedManager"
         reveal_type(Store().purchases.queryset_method())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet"
-        reveal_type(Store().purchases.filter())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet[myapp.models.purchase.Purchase]"
+        reveal_type(Store().purchases.filter())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet"
         reveal_type(Store().purchases.filter().queryset_method())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet"
         reveal_type(User().purchases)  # N: Revealed type is "myapp.models.purchase.Purchase_RelatedManager"
         reveal_type(User().purchases.queryset_method())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet"
-        reveal_type(User().purchases.filter())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet[myapp.models.purchase.Purchase]"
+        reveal_type(User().purchases.filter())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet"
         reveal_type(User().purchases.filter().queryset_method())  # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet"
     installed_apps:
         - myapp

tests/typecheck/managers/querysets/test_as_manager.yml

@@ -310,7 +310,7 @@
 -   case: handles_type_collision_with_from_queryset
     main: |
         from myapp.models import MyModel, FromQuerySet
-        reveal_type(FromQuerySet)  # N: Revealed type is "def [_T <: django.db.models.base.Model] () -> myapp.models.ManagerFromModelQuerySet[_T`1]"
+        reveal_type(FromQuerySet)  # N: Revealed type is "def () -> myapp.models.ManagerFromModelQuerySet[_T`1]"
         reveal_type(MyModel.from_queryset)  # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]"
         reveal_type(MyModel.as_manager)  # N: Revealed type is "myapp.models.ManagerFromModelQuerySet[myapp.models.MyModel]"
     installed_apps:

tests/typecheck/managers/querysets/test_from_queryset.yml

@@ -158,7 +158,7 @@
 -   case: from_queryset_annotates_manager_variable_as_type
     main: |
         from myapp.models import NewManager
-        reveal_type(NewManager)  # N: Revealed type is "def [_T <: django.db.models.base.Model] () -> myapp.models.ManagerFromModelQuerySet[_T`1]"
+        reveal_type(NewManager)  # N: Revealed type is "def () -> myapp.models.ManagerFromModelQuerySet[_T`1]"
     installed_apps:
         - myapp
     files:
@@ -197,7 +197,7 @@
 -   case: from_queryset_returns_intersection_of_manager_and_queryset
     main: |
         from myapp.models import MyModel, NewManager
-        reveal_type(NewManager())  # N: Revealed type is "myapp.models.ModelBaseManagerFromModelQuerySet[<nothing>]"
+        reveal_type(NewManager())  # N: Revealed type is "myapp.models.ModelBaseManagerFromModelQuerySet[_T`1]"
         reveal_type(MyModel.objects)  # N: Revealed type is "myapp.models.ModelBaseManagerFromModelQuerySet[myapp.models.MyModel]"
         reveal_type(MyModel.objects.get())  # N: Revealed type is "myapp.models.MyModel"
         reveal_type(MyModel.objects.manager_only_method())  # N: Revealed type is "builtins.int"
@@ -228,7 +228,7 @@
         reveal_type(MyModel.objects.get())  # N: Revealed type is "myapp.models.MyModel"
         reveal_type(MyModel.objects.manager_only_method())  # N: Revealed type is "builtins.int"
         reveal_type(MyModel.objects.manager_and_queryset_method())  # N: Revealed type is "builtins.str"
-        reveal_type(OtherManager())  # N: Revealed type is "myapp.models.X[<nothing>]"
+        reveal_type(OtherManager())  # N: Revealed type is "myapp.models.X[_T`1]"
         reveal_type(OtherModel.objects)  # N: Revealed type is "myapp.models.X[myapp.models.OtherModel]"
         reveal_type(OtherModel.objects.manager_only_method())  # N: Revealed type is "builtins.int"
         reveal_type(OtherModel.objects.manager_and_queryset_method())  # N: Revealed type is "builtins.str"
@@ -572,8 +572,8 @@
 -   case: reuses_type_when_called_twice_identically
     main: |
         from myapp.models import MyModel, FirstManager, SecondManager
-        reveal_type(FirstManager)  # N: Revealed type is "def [_T <: django.db.models.base.Model] () -> myapp.models.BaseManagerFromModelQuerySet[_T`1]"
-        reveal_type(SecondManager)  # N: Revealed type is "def [_T <: django.db.models.base.Model] () -> myapp.models.BaseManagerFromModelQuerySet[_T`1]"
+        reveal_type(FirstManager)  # N: Revealed type is "def () -> myapp.models.BaseManagerFromModelQuerySet[_T`1]"
+        reveal_type(SecondManager)  # N: Revealed type is "def () -> myapp.models.BaseManagerFromModelQuerySet[_T`1]"
         reveal_type(MyModel.first)  # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet[myapp.models.MyModel]"
         reveal_type(MyModel.second)  # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet[myapp.models.MyModel]"
     installed_apps:
@@ -619,7 +619,7 @@
     main: |
         from myapp.models import MyModel, Generated, BaseManagerFromModelQuerySet
         reveal_type(BaseManagerFromModelQuerySet)  # N: Revealed type is "builtins.int"
-        reveal_type(Generated())  # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet1[<nothing>]"
+        reveal_type(Generated())  # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet1[_T`1]"
         reveal_type(MyModel.objects)  # N: Revealed type is "myapp.models.BaseManagerFromModelQuerySet1[myapp.models.MyModel]"
     installed_apps:
         - myapp
@@ -641,8 +641,8 @@
 -   case: accepts_explicit_none_as_class_name
     main: |
         from myapp.models import PositionalNone, NoneAsKwarg
-        reveal_type(PositionalNone)  # N: Revealed type is "def [_T <: django.db.models.base.Model] () -> myapp.models.BaseManagerFromModelQuerySet[_T`1]"
-        reveal_type(NoneAsKwarg)  # N: Revealed type is "def [_T <: django.db.models.base.Model] () -> myapp.models.BaseManagerFromModelQuerySet[_T`1]"
+        reveal_type(PositionalNone)  # N: Revealed type is "def () -> myapp.models.BaseManagerFromModelQuerySet[_T`1]"
+        reveal_type(NoneAsKwarg)  # N: Revealed type is "def () -> myapp.models.BaseManagerFromModelQuerySet[_T`1]"
     installed_apps:
         - myapp
     files:
@@ -661,7 +661,7 @@
 -   case: uses_fallback_class_name_when_argument_is_not_string_expression
     main: |
         from myapp.models import StrCallable
-        reveal_type(StrCallable)  # N: Revealed type is "def [_T <: django.db.models.base.Model] () -> myapp.models.BaseManagerFromModelQuerySet[_T`1]"
+        reveal_type(StrCallable)  # N: Revealed type is "def () -> myapp.models.BaseManagerFromModelQuerySet[_T`1]"
     installed_apps:
         - myapp
     files:
@@ -675,3 +675,34 @@
                   ...
 
               StrCallable = BaseManager.from_queryset(ModelQuerySet, class_name=str(1))
+
+-   case: test_type_generated_from_queryset_can_be_inherited
+    main: |
+        from myapp.models import MyModel
+        reveal_type(MyModel.objects)
+        reveal_type(MyModel.objects.my_method())
+        reveal_type(MyModel.objects.queryset_method())
+    out: |
+        main:2: note: Revealed type is "myapp.models.MyModelManager[myapp.models.MyModel]"
+        main:3: note: Revealed type is "builtins.str"
+        main:4: note: Revealed type is "builtins.int"
+    installed_apps:
+        - myapp
+    files:
+        - path: myapp/__init__.py
+        - path: myapp/models.py
+          content: |
+              from __future__ import annotations
+              from django.db import models
+
+              class CustomQuerySet(models.QuerySet["MyModel"]):
+                  def queryset_method(self) -> int:
+                      ...
+
+              GeneratedManager = models.Manager.from_queryset(CustomQuerySet)
+
+              class MyModelManager(GeneratedManager):
+                  def my_method(self) -> str:
+                      ...
+              class MyModel(models.Model):
+                  objects = MyModelManager()