Reparametrize managers without explicit type parameters

This extracts the reparametrization logic from typeddjango#1030 in addition to
removing the codepath that copied methods from querysets to managers.
That code path seems to not be needed with this change.

Author: ljodal

django-stubs/contrib/sessions/base_session.pyi

@@ -1,12 +1,14 @@
 from datetime import datetime
-from typing import Any, Dict, Optional, Type
+from typing import Any, Dict, Optional, Type, TypeVar
 
 from django.contrib.sessions.backends.base import SessionBase
 from django.db import models
 
-class BaseSessionManager(models.Manager):
+_T = TypeVar("_T", bound="AbstractBaseSession")
+
+class BaseSessionManager(models.Manager[_T]):
     def encode(self, session_dict: Dict[str, int]) -> str: ...
-    def save(self, session_key: str, session_dict: Dict[str, int], expire_date: datetime) -> AbstractBaseSession: ...
+    def save(self, session_key: str, session_dict: Dict[str, int], expire_date: datetime) -> _T: ...
 
 class AbstractBaseSession(models.Model):
     expire_date: datetime

django-stubs/contrib/sessions/models.pyi

@@ -1,4 +1,8 @@
+from typing import TypeVar
+
 from django.contrib.sessions.base_session import AbstractBaseSession, BaseSessionManager
 
-class SessionManager(BaseSessionManager): ...
+_T = TypeVar("_T", bound="Session")
+
+class SessionManager(BaseSessionManager[_T]): ...
 class Session(AbstractBaseSession): ...

django-stubs/contrib/sites/managers.pyi

@@ -1,6 +1,9 @@
-from typing import Optional
+from typing import Optional, TypeVar
 
+from django.contrib.sites.models import Site
 from django.db import models
 
-class CurrentSiteManager(models.Manager):
+_T = TypeVar("_T", bound=Site)
+
+class CurrentSiteManager(models.Manager[_T]):
     def __init__(self, field_name: Optional[str] = ...) -> None: ...

mypy_django_plugin/lib/helpers.py

@@ -1,5 +1,5 @@
 from collections import OrderedDict
-from typing import TYPE_CHECKING, Any, Dict, Iterable, Iterator, List, Optional, Set, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, Iterable, Iterator, List, Optional, Set, Union
 
 from django.db.models.fields import Field
 from django.db.models.fields.related import RelatedField
@@ -10,11 +10,9 @@
 from mypy.nodes import (
     GDEF,
     MDEF,
-    Argument,
     Block,
     ClassDef,
     Expression,
-    FuncDef,
     MemberExpr,
     MypyFile,
     NameExpr,
@@ -34,11 +32,10 @@
     MethodContext,
     SemanticAnalyzerPluginInterface,
 )
-from mypy.plugins.common import add_method_to_class
 from mypy.semanal import SemanticAnalyzer
-from mypy.types import AnyType, CallableType, Instance, NoneTyp, TupleType
+from mypy.types import AnyType, Instance, NoneTyp, TupleType
 from mypy.types import Type as MypyType
-from mypy.types import TypedDictType, TypeOfAny, UnboundType, UnionType
+from mypy.types import TypedDictType, TypeOfAny, UnionType
 
 from mypy_django_plugin.lib import fullnames
 from mypy_django_plugin.lib.fullnames import WITH_ANNOTATIONS_FULLNAME
@@ -361,86 +358,6 @@ def add_new_sym_for_info(info: TypeInfo, *, name: str, sym_type: MypyType, no_se
     info.names[name] = SymbolTableNode(MDEF, var, plugin_generated=True, no_serialize=no_serialize)
 
 
-def build_unannotated_method_args(method_node: FuncDef) -> Tuple[List[Argument], MypyType]:
-    prepared_arguments = []
-    try:
-        arguments = method_node.arguments[1:]
-    except AttributeError:
-        arguments = []
-    for argument in arguments:
-        argument.type_annotation = AnyType(TypeOfAny.unannotated)
-        prepared_arguments.append(argument)
-    return_type = AnyType(TypeOfAny.unannotated)
-    return prepared_arguments, return_type
-
-
-def bind_or_analyze_type(t: MypyType, api: SemanticAnalyzer, module_name: Optional[str] = None) -> Optional[MypyType]:
-    """Analyze a type. If an unbound type, try to look it up in the given module name.
-
-    That should hopefully give a bound type."""
-    if isinstance(t, UnboundType) and module_name is not None:
-        node = api.lookup_fully_qualified_or_none(module_name + "." + t.name)
-        if node is not None and node.type is not None:
-            return node.type
-
-    return api.anal_type(t)
-
-
-def copy_method_to_another_class(
-    api: SemanticAnalyzer,
-    cls: ClassDef,
-    self_type: Instance,
-    new_method_name: str,
-    method_node: FuncDef,
-    return_type: Optional[MypyType] = None,
-    original_module_name: Optional[str] = None,
-) -> bool:
-    if method_node.type is None:
-        arguments, return_type = build_unannotated_method_args(method_node)
-        add_method_to_class(api, cls, new_method_name, args=arguments, return_type=return_type, self_type=self_type)
-        return True
-
-    method_type = method_node.type
-    if not isinstance(method_type, CallableType):
-        if not api.final_iteration:
-            api.defer()
-        return False
-
-    if return_type is None:
-        return_type = bind_or_analyze_type(method_type.ret_type, api, original_module_name)
-    if return_type is None:
-        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
-    # a method loaded from cache. As mypy doesn't serialize `FuncDef.arguments` when caching
-    arguments = []
-    # Note that the first argument is excluded, as that's `self`
-    for pos, (arg_type, arg_kind, arg_name) in enumerate(
-        zip(method_type.arg_types[1:], method_type.arg_kinds[1:], method_type.arg_names[1:]),
-        start=1,
-    ):
-        bound_arg_type = bind_or_analyze_type(arg_type, api, original_module_name)
-        if bound_arg_type is None:
-            return False
-        if arg_name is None and hasattr(method_node, "arguments"):
-            arg_name = method_node.arguments[pos].variable.name
-        arguments.append(
-            Argument(
-                # Positional only arguments can have name as `None`, if we can't find a name, we just invent one..
-                variable=Var(name=arg_name if arg_name is not None else str(pos), type=arg_type),
-                type_annotation=bound_arg_type,
-                initializer=None,
-                kind=arg_kind,
-                pos_only=arg_name is None,
-            )
-        )
-
-    add_method_to_class(api, 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)
     if sym is not None and isinstance(sym.node, TypeInfo):

mypy_django_plugin/main.py

@@ -26,6 +26,7 @@
 from mypy_django_plugin.transformers.managers import (
     create_new_manager_class_from_as_manager_method,
     create_new_manager_class_from_from_queryset_method,
+    reparametrize_any_manager_hook,
     resolve_manager_method,
 )
 from mypy_django_plugin.transformers.models import (
@@ -240,6 +241,15 @@ def get_method_hook(self, fullname: str) -> Optional[Callable[[MethodContext], M
 
         return None
 
+    def get_customize_class_mro_hook(self, fullname: str) -> Optional[Callable[[ClassDefContext], None]]:
+        sym = self.lookup_fully_qualified(fullname)
+        if (
+            sym is not None
+            and isinstance(sym.node, TypeInfo)
+            and sym.node.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME)
+        ):
+            return reparametrize_any_manager_hook
+
     def get_base_class_hook(self, fullname: str) -> Optional[Callable[[ClassDefContext], None]]:
         # Base class is a Model class definition
         if (

mypy_django_plugin/transformers/managers.py

@@ -15,7 +15,7 @@
     TypeInfo,
     Var,
 )
-from mypy.plugin import AttributeContext, DynamicClassDefContext
+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, Instance, ProperType
@@ -466,3 +466,63 @@ def create_new_manager_class_from_as_manager_method(ctx: DynamicClassDefContext)
         # Note that the generated manager type is always inserted at module level
         SymbolTableNode(GDEF, new_manager_info, plugin_generated=True),
     )
+
+
+def reparametrize_any_manager_hook(ctx: ClassDefContext) -> None:
+    """
+    Add explicit generics to manager classes that are defined without generic.
+
+    Eg.
+
+        class MyManager(models.Manager): ...
+
+    is interpreted as::
+        _T = TypeVar('_T', covariant=True)
+        class MyManager(models.Manager[_T]): ...
+    """
+
+    manager = ctx.api.lookup_fully_qualified_or_none(ctx.cls.fullname)
+    if manager is None or manager.node is None:
+        return
+    assert isinstance(manager.node, TypeInfo)
+
+    if manager.node.type_vars:
+        # We've already been here
+        return
+
+    parent_manager = next(
+        (base for base in manager.node.bases if base.type.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME)),
+        None,
+    )
+    if parent_manager is None:
+        return
+
+    preserve_typevars = (
+        # If args are missing, but tvars present, don't ignore: we have to reparametrize
+        not parent_manager.type.type_vars
+        or parent_manager.args
+        and (
+            not isinstance(parent_manager.args[0], AnyType) or parent_manager.args[0].type_of_any == TypeOfAny.explicit
+        )
+    )
+    if preserve_typevars:
+        return
+
+    base_manager = ctx.api.lookup_fully_qualified_or_none(fullnames.BASE_MANAGER_CLASS_FULLNAME)
+    if base_manager is None:
+        if not ctx.api.final_iteration:
+            ctx.api.defer()
+        return
+    assert isinstance(base_manager.node, TypeInfo)
+
+    tvars = tuple(base_manager.node.defn.type_vars)
+    # For some reason, we have to defer now, otherwise `defer` is called in other place
+    # on final iteration (`SemanticAnalyzer.analyze_func_def`).
+    if any(has_placeholder(tvar) for tvar in tvars):
+        assert not ctx.api.final_iteration, "Too late to reparametrize"
+        ctx.api.defer()
+
+    parent_manager.args = tvars
+    manager.node.type_vars = []
+    manager.node.defn.type_vars = list(tvars)
+    manager.node.add_type_vars()

mypy_django_plugin/transformers/models.py

@@ -282,45 +282,6 @@ def has_any_parametrized_manager_as_base(self, info: TypeInfo) -> bool:
     def is_any_parametrized_manager(self, typ: Instance) -> bool:
         return typ.type.fullname in fullnames.MANAGER_CLASSES and isinstance(typ.args[0], AnyType)
 
-    def create_new_model_parametrized_manager(self, name: str, base_manager_info: TypeInfo) -> Instance:
-        bases = []
-        for original_base in base_manager_info.bases:
-            if self.is_any_parametrized_manager(original_base):
-                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
-        custom_manager_type = Instance(new_manager_info, [Instance(self.model_classdef.info, [])])
-
-        for name, sym in base_manager_info.names.items():
-            # replace self type with new class, if copying method
-            if isinstance(sym.node, FuncDef):
-                copied_method = helpers.copy_method_to_another_class(
-                    api=self.api,
-                    cls=new_manager_info.defn,
-                    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()
-            if isinstance(new_sym.node, Var):
-                new_var = Var(name, type=sym.type)
-                new_var.info = new_manager_info
-                new_var._fullname = new_manager_info.fullname + "." + name
-                new_sym.node = new_var
-            new_manager_info.names[name] = new_sym
-
-        return custom_manager_type
-
     def lookup_manager(self, fullname: str, manager: "Manager[Any]") -> Optional[TypeInfo]:
         manager_info = self.lookup_typeinfo(fullname)
         if manager_info is None:
@@ -354,33 +315,17 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
                 # Manager is already typed -> do nothing unless it's a dynamically generated manager
                 self.reparametrize_dynamically_created_manager(manager_name, manager_info)
                 continue
-            elif manager_info is None:
+
+            if manager_info is None:
                 # We couldn't find a manager type, see if we should create one
                 manager_info = self.create_manager_from_from_queryset(manager_name)
 
             if manager_info is None:
                 incomplete_manager_defs.add(manager_name)
                 continue
 
-            if manager_name not in self.model_classdef.info.names or self.is_manager_dynamically_generated(
-                manager_info
-            ):
-                manager_type = Instance(manager_info, [Instance(self.model_classdef.info, [])])
-                self.add_new_node_to_model_class(manager_name, manager_type)
-            elif self.has_any_parametrized_manager_as_base(manager_info):
-                # Ending up here could for instance be due to having a custom _Manager_
-                # that is not built from a custom QuerySet. Another example is a
-                # related manager.
-                manager_class_name = manager.__class__.__name__
-                custom_model_manager_name = manager.model.__name__ + "_" + manager_class_name
-                try:
-                    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, manager_type)
+            manager_type = Instance(manager_info, [Instance(self.model_classdef.info, [])])
+            self.add_new_node_to_model_class(manager_name, manager_type)
 
         if incomplete_manager_defs:
             if not self.api.final_iteration:

tests/typecheck/managers/querysets/test_from_queryset.yml

@@ -162,9 +162,9 @@
 -   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"
+        reveal_type(NewManager())  # N: Revealed type is "myapp.models.ModelBaseManagerFromModelQuerySet[<nothing>]"
         reveal_type(MyModel.objects)  # N: Revealed type is "myapp.models.ModelBaseManagerFromModelQuerySet[myapp.models.MyModel]"
-        reveal_type(MyModel.objects.get())  # N: Revealed type is "Any"
+        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"
     installed_apps:
@@ -188,12 +188,12 @@
 -   case: from_queryset_with_class_name_provided
     main: |
         from myapp.models import MyModel, NewManager, OtherModel, OtherManager
-        reveal_type(NewManager())  # N: Revealed type is "myapp.models.NewManager"
+        reveal_type(NewManager())  # N: Revealed type is "myapp.models.NewManager[<nothing>]"
         reveal_type(MyModel.objects)  # N: Revealed type is "myapp.models.NewManager[myapp.models.MyModel]"
-        reveal_type(MyModel.objects.get())  # N: Revealed type is "Any"
+        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"
+        reveal_type(OtherManager())  # N: Revealed type is "myapp.models.X[<nothing>]"
         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"