WIP: Implement Managers with generic QuerySets

README.md

@@ -95,7 +95,7 @@ option to get extra information about the error.
 ### I cannot use QuerySet or Manager with type annotations
 
 You can get a `TypeError: 'type' object is not subscriptable`
-when you will try to use `QuerySet[MyModel]`, `Manager[MyModel]` or some other Django-based Generic types.
+when you will try to use `QuerySet[MyModel]`, `Manager[MyModel, MyQuerySet]` or some other Django-based Generic types.
 
 This happens because these Django classes do not support [`__class_getitem__`](https://www.python.org/dev/peps/pep-0560/#class-getitem) magic method in runtime.
 
@@ -215,10 +215,14 @@ error: Return type "MyModel" of "create" incompatible with return type "_T" in s
 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:
+should declare your manager with your model as the type variable, and the
+QuerySet type.
 
 ```python
-class MyManager(models.Manager["MyModel"]):
+class MyQuerySet(models.QuerySet["MyModel"]):
+  ...
+
+class MyManager(models.Manager["MyModel", MyQuerySet]):
     ...
 ```
 

django-stubs/contrib/admin/models.pyi

@@ -9,7 +9,7 @@ CHANGE: int
 DELETION: int
 ACTION_FLAG_CHOICES: Any
 
-class LogEntryManager(models.Manager[LogEntry]):
+class LogEntryManager(models.Manager[LogEntry, models.QuerySet[LogEntry]]):
     def log_action(
         self,
         user_id: int,

django-stubs/contrib/auth/base_user.pyi

@@ -4,11 +4,13 @@ from django.db import models
 from django.db.models.base import Model
 from django.db.models.expressions import Combinable
 from django.db.models.fields import BooleanField
+from django.db.models.query import QuerySet
 from typing_extensions import Literal
 
 _T = TypeVar("_T", bound=Model)
+_QS = TypeVar("_QS", bound=QuerySet[_T])  # type: ignore
 
-class BaseUserManager(models.Manager[_T]):
+class BaseUserManager(models.Manager[_T, _QS]):
     @classmethod
     def normalize_email(cls, email: str | None) -> str: ...
     def make_random_password(self, length: int = ..., allowed_chars: str = ...) -> str: ...

django-stubs/contrib/auth/models.pyi

@@ -15,7 +15,7 @@ _AnyUser: TypeAlias = Model | AnonymousUser
 
 def update_last_login(sender: type[AbstractBaseUser], user: AbstractBaseUser, **kwargs: Any) -> None: ...
 
-class PermissionManager(models.Manager[Permission]):
+class PermissionManager(models.Manager[Permission, models.QuerySet[Permission]]):
     def get_by_natural_key(self, codename: str, app_label: str, model: str) -> Permission: ...
 
 class Permission(models.Model):
@@ -27,7 +27,7 @@ class Permission(models.Model):
     codename = models.CharField(max_length=100)
     def natural_key(self) -> tuple[str, str, str]: ...
 
-class GroupManager(models.Manager[Group]):
+class GroupManager(models.Manager[Group, models.QuerySet[Group]]):
     def get_by_natural_key(self, name: str) -> Group: ...
 
 class Group(models.Model):
@@ -38,8 +38,9 @@ class Group(models.Model):
     def natural_key(self) -> tuple[str]: ...
 
 _T = TypeVar("_T", bound=Model)
+_QS = TypeVar("_QS", bound=models.QuerySet[_T])  # type: ignore
 
-class UserManager(BaseUserManager[_T]):
+class UserManager(BaseUserManager[_T, _QS]):
     def create_user(
         self, username: str, email: str | None = ..., password: str | None = ..., **extra_fields: Any
     ) -> _T: ...
@@ -97,9 +98,9 @@ class AnonymousUser:
     def set_password(self, raw_password: str) -> None: ...
     def check_password(self, raw_password: str) -> Any: ...
     @property
-    def groups(self) -> EmptyManager[Group]: ...
+    def groups(self) -> EmptyManager[Group, models.QuerySet[Group]]: ...
     @property
-    def user_permissions(self) -> EmptyManager[Permission]: ...
+    def user_permissions(self) -> EmptyManager[Permission, models.QuerySet[Permission]]: ...
     def get_user_permissions(self, obj: _AnyUser | None = ...) -> set[str]: ...
     def get_group_permissions(self, obj: _AnyUser | None = ...) -> set[Any]: ...
     def get_all_permissions(self, obj: _AnyUser | None = ...) -> set[str]: ...

django-stubs/contrib/contenttypes/models.pyi

@@ -4,7 +4,7 @@ from django.db import models
 from django.db.models.base import Model
 from django.db.models.query import QuerySet
 
-class ContentTypeManager(models.Manager[ContentType]):
+class ContentTypeManager(models.Manager[ContentType, models.QuerySet[ContentType]]):
     def get_by_natural_key(self, app_label: str, model: str) -> ContentType: ...
     def get_for_model(self, model: type[Model] | Model, for_concrete_model: bool = ...) -> ContentType: ...
     def get_for_models(self, *models: Any, for_concrete_models: bool = ...) -> dict[type[Model], ContentType]: ...

django-stubs/contrib/sessions/base_session.pyi

@@ -6,7 +6,7 @@ from django.db import models
 
 _T = TypeVar("_T", bound=AbstractBaseSession)
 
-class BaseSessionManager(models.Manager[_T]):
+class BaseSessionManager(models.Manager[_T, models.QuerySet[_T]]):
     def encode(self, session_dict: dict[str, Any]) -> str: ...
     def save(self, session_key: str, session_dict: dict[str, Any], expire_date: datetime) -> _T: ...
 

django-stubs/contrib/sites/managers.pyi

@@ -5,5 +5,5 @@ from django.db import models
 
 _T = TypeVar("_T", bound=Site)
 
-class CurrentSiteManager(models.Manager[_T]):
+class CurrentSiteManager(models.Manager[_T, models.QuerySet[_T]]):
     def __init__(self, field_name: str | None = ...) -> None: ...

django-stubs/contrib/sites/models.pyi

@@ -5,8 +5,7 @@ from django.http.request import HttpRequest
 
 SITE_CACHE: Any
 
-class SiteManager(models.Manager[Site]):
-    def get_current(self, request: HttpRequest | None = ...) -> Site: ...
+class SiteManager(models.Manager[Site, models.QuerySet[Site]]):
     def clear_cache(self) -> None: ...
     def get_by_natural_key(self, domain: str) -> Site: ...
 

django-stubs/db/models/base.pyi

@@ -7,6 +7,7 @@ from django.core.exceptions import MultipleObjectsReturned as BaseMultipleObject
 from django.core.exceptions import ObjectDoesNotExist, ValidationError
 from django.db.models.manager import BaseManager
 from django.db.models.options import Options
+from django.db.models.query import QuerySet
 
 _Self = TypeVar("_Self", bound=Model)
 
@@ -19,11 +20,11 @@ class ModelState:
 
 class ModelBase(type):
     @property
-    def objects(cls: type[_Self]) -> BaseManager[_Self]: ...  # type: ignore[misc]
+    def objects(cls: type[_Self]) -> BaseManager[_Self, QuerySet[_Self]]: ...  # type: ignore[misc]
     @property
-    def _default_manager(cls: type[_Self]) -> BaseManager[_Self]: ...  # type: ignore[misc]
+    def _default_manager(cls: type[_Self]) -> BaseManager[_Self, QuerySet[_Self]]: ...  # type: ignore[misc]
     @property
-    def _base_manager(cls: type[_Self]) -> BaseManager[_Self]: ...  # type: ignore[misc]
+    def _base_manager(cls: type[_Self]) -> BaseManager[_Self, QuerySet[_Self]]: ...  # type: ignore[misc]
 
 class Model(metaclass=ModelBase):
     class DoesNotExist(ObjectDoesNotExist): ...

django-stubs/db/models/fields/related.pyi

@@ -20,6 +20,7 @@ from django.db.models.fields.reverse_related import ManyToManyRel as ManyToManyR
 from django.db.models.fields.reverse_related import ManyToOneRel as ManyToOneRel
 from django.db.models.fields.reverse_related import OneToOneRel as OneToOneRel
 from django.db.models.manager import RelatedManager
+from django.db.models.query import QuerySet
 from django.db.models.query_utils import FilteredRelation, PathInfo, Q
 from django.utils.functional import _StrOrPromise
 from typing_extensions import Literal
@@ -204,7 +205,7 @@ class OneToOneField(ForeignKey[_ST, _GT]):
 
 class ManyToManyField(RelatedField[_ST, _GT]):
     _pyi_private_set_type: Sequence[Any]
-    _pyi_private_get_type: RelatedManager[Any]
+    _pyi_private_get_type: RelatedManager[Any, QuerySet[Any]]
 
     description: str
     has_null_arg: bool

django-stubs/db/models/manager.pyi

@@ -10,8 +10,9 @@ from django.db.models.query import QuerySet, RawQuerySet
 from django_stubs_ext import ValuesQuerySet
 
 _T = TypeVar("_T", bound=Model, covariant=True)
+_QS = TypeVar("_QS", bound=QuerySet[_T], covariant=True)  # type: ignore
 
-class BaseManager(Generic[_T]):
+class BaseManager(Generic[_T, _QS]):
     creation_counter: int
     auto_created: bool
     use_in_migrations: bool
@@ -24,14 +25,14 @@ class BaseManager(Generic[_T]):
     ) -> tuple[bool, str | None, str | None, tuple[Any, ...] | None, dict[str, Any] | None]: ...
     def check(self, **kwargs: Any) -> list[Any]: ...
     @classmethod
-    def from_queryset(cls, queryset_class: type[QuerySet], class_name: str | None = ...) -> Any: ...
+    def from_queryset(cls, queryset_class: type[_QS], class_name: str | None = ...) -> type[BaseManager[_T, _QS]]: ...
     @classmethod
     def _get_queryset_methods(cls, queryset_class: type) -> dict[str, Any]: ...
     def contribute_to_class(self, cls: type[Model], name: str) -> None: ...
     def db_manager(self: Self, using: str | None = ..., hints: dict[str, Model] | None = ...) -> Self: ...
     @property
     def db(self) -> str: ...
-    def get_queryset(self) -> QuerySet[_T]: ...
+    def get_queryset(self) -> _QS: ...
     # NOTE: The following methods are in common with QuerySet, but note that the use of QuerySet as a return type
     # rather than a self-type (_QS), since Manager's QuerySet-like methods return QuerySets and not Managers.
     def iterator(self, chunk_size: int = ...) -> Iterator[_T]: ...
@@ -107,25 +108,25 @@ class BaseManager(Generic[_T]):
     def datetimes(
         self, field_name: str, kind: str, order: str = ..., tzinfo: datetime.tzinfo | None = ...
     ) -> ValuesQuerySet[_T, datetime.datetime]: ...
-    def none(self) -> QuerySet[_T]: ...
-    def all(self) -> QuerySet[_T]: ...
-    def filter(self, *args: Any, **kwargs: Any) -> QuerySet[_T]: ...
-    def exclude(self, *args: Any, **kwargs: Any) -> QuerySet[_T]: ...
-    def complex_filter(self, filter_obj: Any) -> QuerySet[_T]: ...
+    def none(self) -> _QS: ...
+    def all(self) -> _QS: ...
+    def filter(self, *args: Any, **kwargs: Any) -> _QS: ...
+    def exclude(self, *args: Any, **kwargs: Any) -> _QS: ...
+    def complex_filter(self, filter_obj: Any) -> _QS: ...
     def count(self) -> int: ...
     async def acount(self) -> int: ...
     def union(self, *other_qs: Any, all: bool = ...) -> QuerySet[_T]: ...
     def intersection(self, *other_qs: Any) -> QuerySet[_T]: ...
     def difference(self, *other_qs: Any) -> QuerySet[_T]: ...
     def select_for_update(
         self, nowait: bool = ..., skip_locked: bool = ..., of: Sequence[str] = ..., no_key: bool = ...
-    ) -> QuerySet[_T]: ...
-    def select_related(self, *fields: Any) -> QuerySet[_T]: ...
-    def prefetch_related(self, *lookups: Any) -> QuerySet[_T]: ...
-    def annotate(self, *args: Any, **kwargs: Any) -> QuerySet[_T]: ...
-    def alias(self, *args: Any, **kwargs: Any) -> QuerySet[_T]: ...
-    def order_by(self, *field_names: Any) -> QuerySet[_T]: ...
-    def distinct(self, *field_names: Any) -> QuerySet[_T]: ...
+    ) -> _QS: ...
+    def select_related(self, *fields: Any) -> _QS: ...
+    def prefetch_related(self, *lookups: Any) -> _QS: ...
+    def annotate(self, *args: Any, **kwargs: Any) -> _QS: ...
+    def alias(self, *args: Any, **kwargs: Any) -> _QS: ...
+    def order_by(self, *field_names: Any) -> _QS: ...
+    def distinct(self, *field_names: Any) -> _QS: ...
     # extra() return type won't be supported any time soon
     def extra(
         self,
@@ -137,22 +138,22 @@ class BaseManager(Generic[_T]):
         select_params: Sequence[Any] | None = ...,
     ) -> QuerySet[Any]: ...
     def reverse(self) -> QuerySet[_T]: ...
-    def defer(self, *fields: Any) -> QuerySet[_T]: ...
-    def only(self, *fields: Any) -> QuerySet[_T]: ...
-    def using(self, alias: str | None) -> QuerySet[_T]: ...
+    def defer(self, *fields: Any) -> _QS: ...
+    def only(self, *fields: Any) -> _QS: ...
+    def using(self, alias: str | None) -> _QS: ...
     @property
     def ordered(self) -> bool: ...
 
-class Manager(BaseManager[_T]): ...
+class Manager(BaseManager[_T, _QS]): ...
 
 # Fake to make ManyToMany work
-class RelatedManager(Manager[_T]):
+class RelatedManager(Manager[_T, _QS]):
     related_val: tuple[int, ...]
     def add(self, *objs: _T | int, bulk: bool = ...) -> None: ...
     def remove(self, *objs: _T | int, bulk: bool = ...) -> None: ...
     def set(self, objs: QuerySet[_T] | Iterable[_T | int], *, bulk: bool = ..., clear: bool = ...) -> None: ...
     def clear(self) -> None: ...
-    def __call__(self, *, manager: str) -> RelatedManager[_T]: ...
+    def __call__(self, *, manager: str) -> RelatedManager[_T, _QS]: ...
 
 class ManagerDescriptor:
     manager: BaseManager
@@ -162,5 +163,5 @@ class ManagerDescriptor:
     @overload
     def __get__(self, instance: Model, cls: type[Model] | None = ...) -> NoReturn: ...
 
-class EmptyManager(Manager[_T]):
+class EmptyManager(Manager[_T, _QS]):
     def __init__(self, model: type[_T]) -> None: ...

django-stubs/db/models/query.pyi

@@ -55,7 +55,7 @@ class _QuerySet(Generic[_T, _Row], Collection[_Row], Reversible[_Row], Sized):
         hints: dict[str, Model] | None = ...,
     ) -> None: ...
     @classmethod
-    def as_manager(cls) -> Manager[Any]: ...
+    def as_manager(cls: type[_QS]) -> Manager[_T, _QS]: ...
     def __len__(self) -> int: ...
     def __bool__(self) -> bool: ...
     def __class_getitem__(cls: type[_QS], item: type[_T]) -> type[_QS]: ...

django-stubs/forms/models.pyi

@@ -269,7 +269,7 @@ class ModelChoiceField(ChoiceField):
     to_field_name: str | None
     def __init__(
         self,
-        queryset: None | Manager[models.Model] | QuerySet[models.Model],
+        queryset: None | Manager[models.Model, QuerySet[models.Model]] | QuerySet[models.Model],
         *,
         empty_label: _StrOrPromise | None = ...,
         required: bool = ...,
@@ -302,7 +302,7 @@ class ModelMultipleChoiceField(ModelChoiceField):
     widget: _ClassLevelWidgetT
     hidden_widget: type[Widget]
     default_error_messages: dict[str, str]
-    def __init__(self, queryset: None | Manager[Model] | QuerySet[Model], **kwargs: Any) -> None: ...
+    def __init__(self, queryset: None | Manager[Model, QuerySet[Model]] | QuerySet[Model], **kwargs: Any) -> None: ...
     def to_python(self, value: Any) -> list[Model]: ...  # type: ignore[override]
     def clean(self, value: Any) -> QuerySet[Model]: ...
     def prepare_value(self, value: Any) -> Any: ...

django-stubs/shortcuts.pyi

@@ -36,6 +36,8 @@ def redirect(
 
 _T = TypeVar("_T", bound=Model)
 
-def get_object_or_404(klass: type[_T] | Manager[_T] | QuerySet[_T], *args: Any, **kwargs: Any) -> _T: ...
-def get_list_or_404(klass: type[_T] | Manager[_T] | QuerySet[_T], *args: Any, **kwargs: Any) -> list[_T]: ...
+def get_object_or_404(klass: type[_T] | Manager[_T, QuerySet[_T]] | QuerySet[_T], *args: Any, **kwargs: Any) -> _T: ...
+def get_list_or_404(
+    klass: type[_T] | Manager[_T, QuerySet[_T]] | QuerySet[_T], *args: Any, **kwargs: Any
+) -> list[_T]: ...
 def resolve_url(to: Callable | Model | str, *args: Any, **kwargs: Any) -> str: ...

mypy_django_plugin/transformers/managers.py

@@ -448,7 +448,9 @@ def create_new_manager_class_from_as_manager_method(ctx: DynamicClassDefContext)
     # 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.
-    var = Var(name=ctx.name, type=Instance(new_manager_info, [AnyType(TypeOfAny.from_omitted_generics)]))
+    any_type = AnyType(TypeOfAny.from_omitted_generics)
+    manager_type = Instance(new_manager_info, [any_type, any_type])
+    var = Var(name=ctx.name, type=manager_type)
     var.info = new_manager_info
     var._fullname = f"{current_module.fullname}.{ctx.name}"
     var.is_inferred = True
@@ -479,7 +481,7 @@ class MyManager(models.Manager): ...
     is interpreted as:
 
         _T = TypeVar('_T', covariant=True)
-        class MyManager(models.Manager[_T]): ...
+        class MyManager(models.Manager[_T, models.QuerySet[_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.
@@ -501,16 +503,17 @@ class MyManager(models.Manager[_T]): ...
     if parent_manager is None:
         return
 
+    type_vars = tuple(parent_manager.type.defn.type_vars)
+
     is_missing_params = (
-        len(parent_manager.args) == 1
+        len(parent_manager.args) == len(type_vars)
         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):
@@ -519,6 +522,13 @@ class MyManager(models.Manager[_T]): ...
         else:
             return
 
+    # FIXME: This results in the following error
+    #
+    #   error: Name "_T" is not defined
+    #
+    # This error goes away if you define literally anything named _T in the
+    # scope of the class. We need to fix this so the type var is set
+    # appropriately without conflicting with any names.
     parent_manager.args = type_vars
     manager.node.defn.type_vars = list(type_vars)
     manager.node.add_type_vars()