Reparametrize managers without explicit type parameters (#1169)

* Reparametrize managers without explicit type parameters

This extracts the reparametrization logic from #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.

* Use typevars from parent instead of base

* Use typevars from parent manager instead of base manager

This removes warnings when subclassing from something other than the
base manager class, where the typevar has been restricted.

* Remove unused imports

* Fix failed test

* Only reparametrize if generics are omitted

* Fix docstring

* Add test with disallow_any_generics=True

* Add an FAQ section and document disallow_any_generics behaviour

Author: ljodal

README.md

@@ -191,6 +191,36 @@ def use_my_model() -> int:
     return foo.xyz # Gives an error
 ```
 
+### Why am I getting incompatible return type errors on my custom managers?
+
+If you declare your custom managers without generics and override built-in
+methods you might see an error message about incompatible error messages,
+something like this:
+
+```python
+from django.db import models
+
+class MyManager(model.Manager):
+    def create(self, **kwargs) -> "MyModel":
+        pass
+```
+
+will cause this error message:
+
+```
+error: Return type "MyModel" of "create" incompatible with return type "_T" in supertype "BaseManager"
+```
+
+This is happening because the `Manager` class is generic, but without
+specifying generics the built-in manager methods are expected to return the
+generic type of the base manager, which is any model. To fix this issue you
+should declare your manager with your model as the type variable:
+
+```python
+class MyManager(models.Manager["MyModel"]):
+    ...
+```
+
 ### How do I annotate cases where I called QuerySet.annotate?
 
 Django-stubs provides a special type, `django_stubs_ext.WithAnnotations[Model]`, which indicates that the `Model` has

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,59 @@ 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 implicit 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]): ...
+
+    Note that this does not happen if mypy is run with disallow_any_generics = True,
+    as not specifying the generic type is then considered an error.
+    """
+
+    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
+
+    is_missing_params = (
+        len(parent_manager.args) == 1
+        and isinstance(parent_manager.args[0], AnyType)
+        and parent_manager.args[0].type_of_any is TypeOfAny.from_omitted_generics
+    )
+    if not is_missing_params:
+        return
+
+    type_vars = tuple(parent_manager.type.defn.type_vars)
+
+    # If we end up with placeholders we need to defer so the placeholders are
+    # resolved in a future iteration
+    if any(has_placeholder(type_var) for type_var in type_vars):
+        if not ctx.api.final_iteration:
+            ctx.api.defer()
+        else:
+            return
+
+    parent_manager.args = type_vars
+    manager.node.defn.type_vars = list(type_vars)
+    manager.node.add_type_vars()

mypy_django_plugin/transformers/models.py

@@ -5,7 +5,7 @@
 from django.db.models.fields.related import ForeignKey
 from django.db.models.fields.reverse_related import ManyToManyRel, ManyToOneRel, OneToOneRel
 from mypy.checker import TypeChecker
-from mypy.nodes import ARG_STAR2, Argument, AssignmentStmt, CallExpr, Context, FuncDef, NameExpr, TypeInfo, Var
+from mypy.nodes import ARG_STAR2, Argument, AssignmentStmt, CallExpr, Context, NameExpr, TypeInfo, Var
 from mypy.plugin import AnalyzeTypeContext, AttributeContext, CheckerPluginInterface, ClassDefContext
 from mypy.plugins import common
 from mypy.semanal import SemanticAnalyzer
@@ -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: