Make QuerySet[_T] an external alias to _QuerySet[_T, _T].

This currently has the drawback that error messages display the internal type _QuerySet, with both type arguments.

See also discussion on typeddjango#661 and typeddjango#608.

Fixes typeddjango#635: QuerySet methods on Managers (like .all()) now return QuerySets rather than Managers.

Address code review by @sobolevn.

Author: syastrov

django-stubs/db/models/manager.pyi

@@ -1,12 +1,30 @@
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, TypeVar, Union
+import datetime
+from typing import (
+    Any,
+    Dict,
+    Generic,
+    Iterable,
+    Iterator,
+    List,
+    MutableMapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+)
 
+from django.db.models import Combinable
 from django.db.models.base import Model
-from django.db.models.query import QuerySet
+from django.db.models.query import QuerySet, RawQuerySet
+
+from django_stubs_ext import ValuesQuerySet
 
 _T = TypeVar("_T", bound=Model, covariant=True)
 _M = TypeVar("_M", bound="BaseManager")
 
-class BaseManager(QuerySet[_T]):
+class BaseManager(Generic[_T]):
     creation_counter: int = ...
     auto_created: bool = ...
     use_in_migrations: bool = ...
@@ -24,6 +42,80 @@ class BaseManager(QuerySet[_T]):
     def contribute_to_class(self, model: Type[Model], name: str) -> None: ...
     def db_manager(self: _M, using: Optional[str] = ..., hints: Optional[Dict[str, Model]] = ...) -> _M: ...
     def get_queryset(self) -> QuerySet[_T]: ...
+    # 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]: ...
+    def aggregate(self, *args: Any, **kwargs: Any) -> Dict[str, Any]: ...
+    def get(self, *args: Any, **kwargs: Any) -> _T: ...
+    def create(self, *args: Any, **kwargs: Any) -> _T: ...
+    def bulk_create(
+        self, objs: Iterable[_T], batch_size: Optional[int] = ..., ignore_conflicts: bool = ...
+    ) -> List[_T]: ...
+    def bulk_update(self, objs: Iterable[_T], fields: Sequence[str], batch_size: Optional[int] = ...) -> None: ...
+    def get_or_create(self, defaults: Optional[MutableMapping[str, Any]] = ..., **kwargs: Any) -> Tuple[_T, bool]: ...
+    def update_or_create(
+        self, defaults: Optional[MutableMapping[str, Any]] = ..., **kwargs: Any
+    ) -> Tuple[_T, bool]: ...
+    def earliest(self, *fields: Any, field_name: Optional[Any] = ...) -> _T: ...
+    def latest(self, *fields: Any, field_name: Optional[Any] = ...) -> _T: ...
+    def first(self) -> Optional[_T]: ...
+    def last(self) -> Optional[_T]: ...
+    def in_bulk(self, id_list: Iterable[Any] = ..., *, field_name: str = ...) -> Dict[Any, _T]: ...
+    def delete(self) -> Tuple[int, Dict[str, int]]: ...
+    def update(self, **kwargs: Any) -> int: ...
+    def exists(self) -> bool: ...
+    def explain(self, *, format: Optional[Any] = ..., **options: Any) -> str: ...
+    def raw(
+        self,
+        raw_query: str,
+        params: Any = ...,
+        translations: Optional[Dict[str, str]] = ...,
+        using: Optional[str] = ...,
+    ) -> RawQuerySet: ...
+    # The type of values may be overridden to be more specific in the mypy plugin, depending on the fields param
+    def values(self, *fields: Union[str, Combinable], **expressions: Any) -> ValuesQuerySet[_T, Dict[str, Any]]: ...
+    # The type of values_list may be overridden to be more specific in the mypy plugin, depending on the fields param
+    def values_list(
+        self, *fields: Union[str, Combinable], flat: bool = ..., named: bool = ...
+    ) -> ValuesQuerySet[_T, Any]: ...
+    def dates(self, field_name: str, kind: str, order: str = ...) -> ValuesQuerySet[_T, datetime.date]: ...
+    def datetimes(
+        self, field_name: str, kind: str, order: str = ..., tzinfo: Optional[datetime.tzinfo] = ...
+    ) -> 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 count(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]: ...
+    # extra() return type won't be supported any time soon
+    def extra(
+        self,
+        select: Optional[Dict[str, Any]] = ...,
+        where: Optional[List[str]] = ...,
+        params: Optional[List[Any]] = ...,
+        tables: Optional[List[str]] = ...,
+        order_by: Optional[Sequence[str]] = ...,
+        select_params: Optional[Sequence[Any]] = ...,
+    ) -> 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: Optional[str]) -> QuerySet[_T]: ...
+    @property
+    def ordered(self) -> bool: ...
 
 class Manager(BaseManager[_T]): ...
 

django-stubs/db/models/query.pyi

@@ -28,9 +28,10 @@ from django.db.models.query_utils import Q as Q  # noqa: F401
 from django.db.models.sql.query import Query, RawQuery
 
 _T = TypeVar("_T", bound=models.Model, covariant=True)
-_QS = TypeVar("_QS", bound="QuerySet")
+_Row = TypeVar("_Row", covariant=True)
+_QS = TypeVar("_QS", bound="_QuerySet")
 
-class QuerySet(Generic[_T], Collection[_T], Reversible[_T], Sized):
+class _QuerySet(Generic[_T, _Row], Collection[_Row], Reversible[_Row], Sized):
     model: Type[_T]
     query: Query
     def __init__(
@@ -47,11 +48,13 @@ class QuerySet(Generic[_T], Collection[_T], Reversible[_T], Sized):
     def __class_getitem__(cls: Type[_QS], item: Type[_T]) -> Type[_QS]: ...
     def __getstate__(self) -> Dict[str, Any]: ...
     # Technically, the other QuerySet must be of the same type _T, but _T is covariant
-    def __and__(self: _QS, other: QuerySet[_T]) -> _QS: ...
-    def __or__(self: _QS, other: QuerySet[_T]) -> _QS: ...
-    def iterator(self, chunk_size: int = ...) -> Iterator[_T]: ...
+    def __and__(self: _QS, other: _QuerySet[_T, _Row]) -> _QS: ...
+    def __or__(self: _QS, other: _QuerySet[_T, _Row]) -> _QS: ...
+    # IMPORTANT: When updating any of the following methods' signatures, please ALSO modify
+    # the corresponding method in BaseManager.
+    def iterator(self, chunk_size: int = ...) -> Iterator[_Row]: ...
     def aggregate(self, *args: Any, **kwargs: Any) -> Dict[str, Any]: ...
-    def get(self, *args: Any, **kwargs: Any) -> _T: ...
+    def get(self, *args: Any, **kwargs: Any) -> _Row: ...
     def create(self, *args: Any, **kwargs: Any) -> _T: ...
     def bulk_create(
         self, objs: Iterable[_T], batch_size: Optional[int] = ..., ignore_conflicts: bool = ...
@@ -61,10 +64,10 @@ class QuerySet(Generic[_T], Collection[_T], Reversible[_T], Sized):
     def update_or_create(
         self, defaults: Optional[MutableMapping[str, Any]] = ..., **kwargs: Any
     ) -> Tuple[_T, bool]: ...
-    def earliest(self, *fields: Any, field_name: Optional[Any] = ...) -> _T: ...
-    def latest(self, *fields: Any, field_name: Optional[Any] = ...) -> _T: ...
-    def first(self) -> Optional[_T]: ...
-    def last(self) -> Optional[_T]: ...
+    def earliest(self, *fields: Any, field_name: Optional[Any] = ...) -> _Row: ...
+    def latest(self, *fields: Any, field_name: Optional[Any] = ...) -> _Row: ...
+    def first(self) -> Optional[_Row]: ...
+    def last(self) -> Optional[_Row]: ...
     def in_bulk(self, id_list: Iterable[Any] = ..., *, field_name: str = ...) -> Dict[Any, _T]: ...
     def delete(self) -> Tuple[int, Dict[str, int]]: ...
     def update(self, **kwargs: Any) -> int: ...
@@ -78,15 +81,15 @@ class QuerySet(Generic[_T], Collection[_T], Reversible[_T], Sized):
         using: Optional[str] = ...,
     ) -> RawQuerySet: ...
     # The type of values may be overridden to be more specific in the mypy plugin, depending on the fields param
-    def values(self, *fields: Union[str, Combinable], **expressions: Any) -> _ValuesQuerySet[_T, Dict[str, Any]]: ...
+    def values(self, *fields: Union[str, Combinable], **expressions: Any) -> _QuerySet[_T, Dict[str, Any]]: ...
     # The type of values_list may be overridden to be more specific in the mypy plugin, depending on the fields param
     def values_list(
         self, *fields: Union[str, Combinable], flat: bool = ..., named: bool = ...
-    ) -> _ValuesQuerySet[_T, Any]: ...
-    def dates(self, field_name: str, kind: str, order: str = ...) -> _ValuesQuerySet[_T, datetime.date]: ...
+    ) -> _QuerySet[_T, Any]: ...
+    def dates(self, field_name: str, kind: str, order: str = ...) -> _QuerySet[_T, datetime.date]: ...
     def datetimes(
         self, field_name: str, kind: str, order: str = ..., tzinfo: Optional[datetime.tzinfo] = ...
-    ) -> _ValuesQuerySet[_T, datetime.datetime]: ...
+    ) -> _QuerySet[_T, datetime.datetime]: ...
     def none(self: _QS) -> _QS: ...
     def all(self: _QS) -> _QS: ...
     def filter(self: _QS, *args: Any, **kwargs: Any) -> _QS: ...
@@ -114,7 +117,7 @@ class QuerySet(Generic[_T], Collection[_T], Reversible[_T], Sized):
         tables: Optional[List[str]] = ...,
         order_by: Optional[Sequence[str]] = ...,
         select_params: Optional[Sequence[Any]] = ...,
-    ) -> QuerySet[Any]: ...
+    ) -> _QuerySet[Any, Any]: ...
     def reverse(self: _QS) -> _QS: ...
     def defer(self: _QS, *fields: Any) -> _QS: ...
     def only(self: _QS, *fields: Any) -> _QS: ...
@@ -124,28 +127,13 @@ class QuerySet(Generic[_T], Collection[_T], Reversible[_T], Sized):
     @property
     def db(self) -> str: ...
     def resolve_expression(self, *args: Any, **kwargs: Any) -> Any: ...
-    def __iter__(self) -> Iterator[_T]: ...
+    def __iter__(self) -> Iterator[_Row]: ...
     def __contains__(self, x: object) -> bool: ...
     @overload
-    def __getitem__(self, i: int) -> _T: ...
-    @overload
-    def __getitem__(self: _QS, s: slice) -> _QS: ...
-    def __reversed__(self) -> Iterator[_T]: ...
-
-_Row = TypeVar("_Row", covariant=True)
-
-class _ValuesQuerySet(QuerySet[_T], Collection[_Row], Reversible[_Row], Sized):
-    def __iter__(self) -> Iterator[_Row]: ...  # type: ignore
-    @overload  # type: ignore
     def __getitem__(self, i: int) -> _Row: ...
     @overload
     def __getitem__(self: _QS, s: slice) -> _QS: ...
-    def iterator(self, chunk_size: int = ...) -> Iterator[_Row]: ...  # type: ignore
-    def get(self, *args: Any, **kwargs: Any) -> _Row: ...  # type: ignore
-    def earliest(self, *fields: Any, field_name: Optional[Any] = ...) -> _Row: ...  # type: ignore
-    def latest(self, *fields: Any, field_name: Optional[Any] = ...) -> _Row: ...  # type: ignore
-    def first(self) -> Optional[_Row]: ...  # type: ignore
-    def last(self) -> Optional[_Row]: ...  # type: ignore
+    def __reversed__(self) -> Iterator[_Row]: ...
 
 class RawQuerySet(Iterable[_T], Sized):
     query: RawQuery
@@ -179,6 +167,8 @@ class RawQuerySet(Iterable[_T], Sized):
     def resolve_model_init_order(self) -> Tuple[List[str], List[int], List[Tuple[str, int]]]: ...
     def using(self, alias: Optional[str]) -> RawQuerySet[_T]: ...
 
+QuerySet = _QuerySet[_T, _T]
+
 class Prefetch(object):
     def __init__(self, lookup: str, queryset: Optional[QuerySet] = ..., to_attr: Optional[str] = ...) -> None: ...
     def __getstate__(self) -> Dict[str, Any]: ...

django-stubs/views/generic/list.pyi

@@ -6,7 +6,7 @@ from django.db.models.query import QuerySet
 from django.http import HttpRequest, HttpResponse
 from django.views.generic.base import ContextMixin, TemplateResponseMixin, View
 
-T = TypeVar("T", bound=Model)
+T = TypeVar("T", bound=Model, covariant=True)
 
 class MultipleObjectMixin(Generic[T], ContextMixin):
     allow_empty: bool = ...

django_stubs_ext/django_stubs_ext/aliases.py

@@ -1,8 +1,8 @@
 import typing
 
 if typing.TYPE_CHECKING:
-    from django.db.models.query import _T, _Row, _ValuesQuerySet
+    from django.db.models.query import _T, _QuerySet, _Row
 
-    ValuesQuerySet = _ValuesQuerySet[_T, _Row]
+    ValuesQuerySet = _QuerySet[_T, _Row]
 else:
     ValuesQuerySet = typing.Any

mypy_django_plugin/django/context.py

@@ -19,7 +19,6 @@
 from mypy.types import AnyType, Instance
 from mypy.types import Type as MypyType
 from mypy.types import TypeOfAny, UnionType
-from typing_extensions import Final
 
 from mypy_django_plugin.lib import fullnames, helpers
 from mypy_django_plugin.lib.fullnames import WITH_ANNOTATIONS_FULLNAME
@@ -116,10 +115,10 @@ def model_modules(self) -> Dict[str, Set[Type[Model]]]:
 
     def get_model_class_by_fullname(self, fullname: str) -> Optional[Type[Model]]:
         """Returns None if Model is abstract"""
-        ANNOTATED_PREFIX: Final = WITH_ANNOTATIONS_FULLNAME + "["
-        if fullname.startswith(ANNOTATED_PREFIX):
+        annotated_prefix = WITH_ANNOTATIONS_FULLNAME + "["
+        if fullname.startswith(annotated_prefix):
             # For our "annotated models", extract the original model fullname
-            fullname = fullname[len(ANNOTATED_PREFIX) :].rstrip("]")
+            fullname = fullname[len(annotated_prefix) :].rstrip("]")
 
         module, _, model_cls_name = fullname.rpartition(".")
         for model_cls in self.model_modules.get(module, set()):

mypy_django_plugin/lib/fullnames.py

@@ -11,8 +11,7 @@
 MANYTOMANY_FIELD_FULLNAME = "django.db.models.fields.related.ManyToManyField"
 DUMMY_SETTINGS_BASE_CLASS = "django.conf._DjangoConfLazyObject"
 
-QUERYSET_CLASS_FULLNAME = "django.db.models.query.QuerySet"
-VALUES_QUERYSET_CLASS_FULLNAME = "django.db.models.query._ValuesQuerySet"
+QUERYSET_CLASS_FULLNAME = "django.db.models.query._QuerySet"
 BASE_MANAGER_CLASS_FULLNAME = "django.db.models.manager.BaseManager"
 MANAGER_CLASS_FULLNAME = "django.db.models.manager.Manager"
 RELATED_MANAGER_CLASS = "django.db.models.manager.RelatedManager"

mypy_django_plugin/lib/helpers.py

@@ -62,7 +62,7 @@ def is_toml(filename: str) -> bool:
 def lookup_fully_qualified_sym(fullname: str, all_modules: Dict[str, MypyFile]) -> Optional[SymbolTableNode]:
     if "." not in fullname:
         return None
-    if "[" in fullname:
+    if "[" in fullname and "]" in fullname:
         # We sometimes generate fake fullnames like a.b.C[x.y.Z] to provide a better representation to users
         # Make sure that we handle lookups of those types of names correctly if the part inside [] contains "."
         bracket_start = fullname.index("[")

mypy_django_plugin/main.py

@@ -266,27 +266,28 @@ def get_method_hook(self, fullname: str) -> Optional[Callable[[MethodContext], M
             if info and info.has_base(fullnames.FORM_MIXIN_CLASS_FULLNAME):
                 return forms.extract_proper_type_for_get_form
 
+        manager_classes = self._get_current_manager_bases()
+
         if method_name == "values":
             info = self._get_typeinfo_or_none(class_fullname)
-            if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME):
+            if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME) or class_fullname in manager_classes:
                 return partial(querysets.extract_proper_type_queryset_values, django_context=self.django_context)
 
         if method_name == "values_list":
             info = self._get_typeinfo_or_none(class_fullname)
-            if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME):
+            if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME) or class_fullname in manager_classes:
                 return partial(querysets.extract_proper_type_queryset_values_list, django_context=self.django_context)
 
         if method_name == "annotate":
             info = self._get_typeinfo_or_none(class_fullname)
-            if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME):
+            if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME) or class_fullname in manager_classes:
                 return partial(querysets.extract_proper_type_queryset_annotate, django_context=self.django_context)
 
         if method_name == "get_field":
             info = self._get_typeinfo_or_none(class_fullname)
             if info and info.has_base(fullnames.OPTIONS_CLASS_FULLNAME):
                 return partial(meta.return_proper_field_type_from_get_field, django_context=self.django_context)
 
-        manager_classes = self._get_current_manager_bases()
         if class_fullname in manager_classes and method_name == "create":
             return partial(init_create.redefine_and_typecheck_model_create, django_context=self.django_context)
         if class_fullname in manager_classes and method_name in {"filter", "get", "exclude"}:

mypy_django_plugin/transformers/models.py

@@ -197,7 +197,6 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
         for manager_name, manager in model_cls._meta.managers_map.items():
             manager_class_name = manager.__class__.__name__
             manager_fullname = helpers.get_class_fullname(manager.__class__)
-
             try:
                 manager_info = self.lookup_typeinfo_or_incomplete_defn_error(manager_fullname)
             except helpers.IncompleteDefnException as exc: