Allow to run from_queryset() with BaseManager, Manager (#271)

* allow to run from_queryset() with BaseManager, Manager

* fix tests

django-stubs/db/models/base.pyi

@@ -2,7 +2,7 @@ from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, Ty
 
 from django.core.checks.messages import CheckMessage
 from django.core.exceptions import ValidationError
-from django.db.models.manager import Manager
+from django.db.models.manager import BaseManager
 from django.db.models.options import Options
 
 _Self = TypeVar("_Self", bound="Model")
@@ -13,9 +13,9 @@ class Model(metaclass=ModelBase):
     class DoesNotExist(Exception): ...
     class MultipleObjectsReturned(Exception): ...
     class Meta: ...
-    _default_manager: Manager[Model]
     _meta: Options[Any]
-    objects: Manager[Any]
+    _default_manager: BaseManager[Model]
+    objects: BaseManager[Any]
     pk: Any = ...
     def __init__(self: _Self, *args, **kwargs) -> None: ...
     def delete(self, using: Any = ..., keep_parents: bool = ...) -> Tuple[int, Dict[str, int]]: ...

mypy.ini

@@ -1 +1,2 @@
 [mypy]
+warn_unused_ignores = True

mypy_django_plugin/lib/fullnames.py

@@ -23,9 +23,7 @@ FORM_MIXIN_CLASS_FULLNAME = 'django.views.generic.edit.FormMixin'
 
 MANAGER_CLASSES = {
     MANAGER_CLASS_FULLNAME,
-    RELATED_MANAGER_CLASS,
     BASE_MANAGER_CLASS_FULLNAME,
-    # QUERYSET_CLASS_FULLNAME
 }
 
 RELATED_FIELDS_CLASSES = {

mypy_django_plugin/transformers/managers.py

@@ -1,5 +1,5 @@
 from mypy.nodes import (
-    GDEF, FuncDef, MemberExpr, NameExpr, StrExpr, SymbolTableNode, TypeInfo,
+    GDEF, FuncDef, MemberExpr, NameExpr, RefExpr, StrExpr, SymbolTableNode, TypeInfo,
 )
 from mypy.plugin import ClassDefContext, DynamicClassDefContext
 from mypy.types import AnyType, Instance, TypeOfAny
@@ -10,9 +10,11 @@ from mypy_django_plugin.lib import helpers
 def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefContext) -> None:
     semanal_api = helpers.get_semanal_api(ctx)
 
-    assert isinstance(ctx.call.callee, MemberExpr)
+    callee = ctx.call.callee
-    assert isinstance(ctx.call.callee.expr, NameExpr)
+    assert isinstance(callee, MemberExpr)
-    base_manager_info = ctx.call.callee.expr.node
+    assert isinstance(callee.expr, RefExpr)
+
+    base_manager_info = callee.expr.node
     if base_manager_info is None:
         if not semanal_api.final_iteration:
             semanal_api.defer()

mypy_django_plugin/transformers/models.py

@@ -144,7 +144,7 @@ class AddManagers(ModelClassInitializer):
         return False
 
     def is_any_parametrized_manager(self, typ: Instance) -> bool:
-        return typ.type.fullname == fullnames.MANAGER_CLASS_FULLNAME and isinstance(typ.args[0], AnyType)
+        return typ.type.fullname in fullnames.MANAGER_CLASSES and isinstance(typ.args[0], AnyType)
 
     def get_generated_manager_mappings(self, base_manager_fullname: str) -> Dict[str, str]:
         base_manager_info = self.lookup_typeinfo(base_manager_fullname)

scripts/enabled_test_modules.py

@@ -271,7 +271,7 @@ IGNORED_ERRORS = {
         '"Manager[Any]" has no attribute "args"',
         'Dict entry 0 has incompatible type "Any"',
         'Argument 1 to "append" of "list" has incompatible type',
-        'base class "Model" defined the type as "Manager[Any]"',
+        'base class "Model" defined the type as "BaseManager[Any]"',
         'Argument 1 to "RunPython" has incompatible type "str"',
 
     ],

test-data/typecheck/managers/querysets/test_from_queryset.yml

@@ -1,4 +1,48 @@
--   case: test_from_queryset_returns_intersection_of_manager_and_queryset
+-   case: from_queryset_with_base_manager
+    main: |
+        from myapp.models import MyModel
+        reveal_type(MyModel().objects)  # N: Revealed type is 'myapp.models.MyModel_NewManager[myapp.models.MyModel]'
+        reveal_type(MyModel().objects.get())  # N: Revealed type is 'myapp.models.MyModel*'
+        reveal_type(MyModel().objects.queryset_method())  # N: Revealed type is 'builtins.str'
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+                from django.db.models.manager import BaseManager
+
+                class ModelQuerySet(models.QuerySet):
+                    def queryset_method(self) -> str:
+                        return 'hello'
+                NewManager = BaseManager.from_queryset(ModelQuerySet)
+                class MyModel(models.Model):
+                    objects = NewManager()
+
+-   case: from_queryset_with_manager
+    main: |
+        from myapp.models import MyModel
+        reveal_type(MyModel().objects)  # N: Revealed type is 'myapp.models.MyModel_NewManager[myapp.models.MyModel]'
+        reveal_type(MyModel().objects.get())  # N: Revealed type is 'myapp.models.MyModel*'
+        reveal_type(MyModel().objects.queryset_method())  # N: Revealed type is 'builtins.str'
+    installed_apps:
+        - myapp
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+
+                class ModelQuerySet(models.QuerySet):
+                    def queryset_method(self) -> str:
+                        return 'hello'
+
+                NewManager = models.Manager.from_queryset(ModelQuerySet)
+                class MyModel(models.Model):
+                    objects = NewManager()
+
+-   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.NewManager'
@@ -24,7 +68,7 @@
                 class MyModel(models.Model):
                     objects = NewManager()
 
--   case: test_from_queryset_with_class_name_provided
+-   case: from_queryset_with_class_name_provided
     main: |
         from myapp.models import MyModel, NewManager
         reveal_type(NewManager())  # N: Revealed type is 'myapp.models.NewManager'

test-data/typecheck/managers/test_managers.yml

@@ -48,7 +48,7 @@
                 class Base(Generic[_T]):
                     def __init__(self, model_cls: Type[_T]):
                         self.model_cls = model_cls
-                        reveal_type(self.model_cls._default_manager)  # N: Revealed type is 'django.db.models.manager.Manager[django.db.models.base.Model]'
+                        reveal_type(self.model_cls._default_manager)  # N: Revealed type is 'django.db.models.manager.BaseManager[django.db.models.base.Model]'
                 class MyModel(models.Model):
                     pass
                 class Child(Base[MyModel]):