make ValuesQuerySet have proper Collection generic type (#140)

2025-12-08 13:04:47 +08:00 · 2019-08-24 18:24:21 +03:00
parent 5fc39ff110
commit fc9a335dfd
11 changed files with 93 additions and 88 deletions
--- a/django-stubs/db/models/expressions.pyi
+++ b/django-stubs/db/models/expressions.pyi
@@ -6,6 +6,7 @@ from django.db.models.sql.compiler import SQLCompiler

 from django.db.models import Q, QuerySet
 from django.db.models.fields import Field
+from django.db.models.query import _BaseQuerySet

 _OutputField = Union[Field, str]

@@ -125,7 +126,7 @@ class Subquery(Expression):
    template: str = ...
    queryset: QuerySet = ...
    extra: Dict[Any, Any] = ...
-    def __init__(self, queryset: QuerySet, output_field: Optional[_OutputField] = ..., **extra: Any) -> None: ...
+    def __init__(self, queryset: _BaseQuerySet, output_field: Optional[_OutputField] = ..., **extra: Any) -> None: ...

 class Exists(Subquery):
    negated: bool = ...
--- a/django-stubs/db/models/query.pyi
+++ b/django-stubs/db/models/query.pyi
@@ -28,9 +28,9 @@ from django.db.models import Manager
 from django.db.models.query_utils import Q as Q

 _T = TypeVar("_T", bound=models.Model, covariant=True)
-_QS = TypeVar("_QS", bound="QuerySet")
+_QS = TypeVar("_QS", bound="_BaseQuerySet")

-class QuerySet(Generic[_T], Collection[_T], Sized):
+class _BaseQuerySet(Generic[_T], Sized):
    query: Query
    def __init__(
        self,
@@ -42,21 +42,13 @@ class QuerySet(Generic[_T], Collection[_T], Sized):
    @classmethod
    def as_manager(cls) -> Manager[Any]: ...
    def __len__(self) -> int: ...
-    def __iter__(self) -> Iterator[_T]: ...
-    def __contains__(self, x: object) -> bool: ...
-    @overload
-    def __getitem__(self, i: int) -> _T: ...
-    @overload
-    def __getitem__(self, s: slice) -> QuerySet[_T]: ...
    def __bool__(self) -> bool: ...
    def __class_getitem__(cls, item: Type[_T]):
        pass
    def __getstate__(self) -> Dict[str, Any]: ...
-    # __and__ and __or__ ignore the other QuerySet's _Row type parameter
-    # because they use the same row type as the self QuerySet.
    # Technically, the other QuerySet must be of the same type _T, but _T is covariant
-    def __and__(self, other: QuerySet[_T]) -> QuerySet[_T]: ...
-    def __or__(self, other: QuerySet[_T]) -> QuerySet[_T]: ...
+    def __and__(self: _QS, other: _BaseQuerySet[_T]) -> _QS: ...
+    def __or__(self: _QS, other: _BaseQuerySet[_T]) -> _QS: ...
    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: ...
@@ -94,22 +86,22 @@ class QuerySet(Generic[_T], Collection[_T], Sized):
    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 none(self: _QS) -> _QS: ...
+    def all(self: _QS) -> _QS: ...
+    def filter(self: _QS, *args: Any, **kwargs: Any) -> _QS: ...
+    def exclude(self: _QS, *args: Any, **kwargs: Any) -> _QS: ...
+    def complex_filter(self, filter_obj: Any) -> _QS: ...
    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: Tuple = ...) -> QuerySet[_T]: ...
-    def select_related(self, *fields: Any) -> QuerySet[_T]: ...
-    def prefetch_related(self, *lookups: Any) -> QuerySet[_T]: ...
+    def union(self: _QS, *other_qs: Any, all: bool = ...) -> _QS: ...
+    def intersection(self: _QS, *other_qs: Any) -> _QS: ...
+    def difference(self: _QS, *other_qs: Any) -> _QS: ...
+    def select_for_update(self: _QS, nowait: bool = ..., skip_locked: bool = ..., of: Tuple = ...) -> _QS: ...
+    def select_related(self: _QS, *fields: Any) -> _QS: ...
+    def prefetch_related(self: _QS, *lookups: Any) -> _QS: ...
    # TODO: return type
    def annotate(self, *args: Any, **kwargs: Any) -> QuerySet[Any]: ...
-    def order_by(self, *field_names: Any) -> QuerySet[_T]: ...
-    def distinct(self, *field_names: Any) -> QuerySet[_T]: ...
+    def order_by(self: _QS, *field_names: Any) -> _QS: ...
+    def distinct(self: _QS, *field_names: Any) -> _QS: ...
    # extra() return type won't be supported any time soon
    def extra(
        self,
@@ -120,10 +112,10 @@ class QuerySet(Generic[_T], Collection[_T], Sized):
        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]: ...
+    def reverse(self: _QS) -> _QS: ...
+    def defer(self: _QS, *fields: Any) -> _QS: ...
+    def only(self: _QS, *fields: Any) -> _QS: ...
+    def using(self: _QS, alias: Optional[str]) -> _QS: ...
    @property
    def ordered(self) -> bool: ...
    @property
@@ -132,10 +124,18 @@ class QuerySet(Generic[_T], Collection[_T], Sized):
    # TODO: remove when django adds __class_getitem__ methods
    def __getattr__(self, item: str) -> Any: ...

+class QuerySet(_BaseQuerySet[_T], Collection[_T], Sized):
+    def __iter__(self) -> Iterator[_T]: ...
+    def __contains__(self, x: object) -> bool: ...
+    @overload
+    def __getitem__(self, i: int) -> _T: ...
+    @overload
+    def __getitem__(self: _QS, s: slice) -> _QS: ...
+
 _Row = TypeVar("_Row", covariant=True)

 class BaseIterable(Sequence[_Row]):
-    def __init__(self, queryset: QuerySet, chunked_fetch: bool = ..., chunk_size: int = ...): ...
+    def __init__(self, queryset: _BaseQuerySet, chunked_fetch: bool = ..., chunk_size: int = ...): ...
    def __iter__(self) -> Iterator[_Row]: ...
    def __contains__(self, x: object) -> bool: ...
    def __len__(self) -> int: ...
@@ -152,53 +152,19 @@ class NamedValuesListIterable(ValuesListIterable): ...
 class FlatValuesListIterable(BaseIterable):
    def __iter__(self) -> Iterator[Any]: ...

-class ValuesQuerySet(Generic[_T, _Row], QuerySet[_T], Collection[_Row], Sized):
+class ValuesQuerySet(_BaseQuerySet[_T], Collection[_Row], Sized):
+    def __contains__(self, x: object) -> bool: ...
    def __iter__(self) -> Iterator[_Row]: ...  # type: ignore
    @overload  # type: ignore
    def __getitem__(self, i: int) -> _Row: ...
    @overload
-    def __getitem__(self, s: slice) -> ValuesQuerySet[_T, _Row]: ...
-    # Technically, the other QuerySet must be of the same type _T, but _T is covariant
-    def __and__(self, other: ValuesQuerySet[_T, _Row]) -> ValuesQuerySet[_T, _Row]: ...  # type: ignore
-    def __or__(self, other: ValuesQuerySet[_T, _Row]) -> ValuesQuerySet[_T, _Row]: ...  # type: ignore
+    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 none(self) -> ValuesQuerySet[_T, _Row]: ...
-    def all(self) -> ValuesQuerySet[_T, _Row]: ...
-    def filter(self, *args: Any, **kwargs: Any) -> ValuesQuerySet[_T, _Row]: ...
-    def exclude(self, *args: Any, **kwargs: Any) -> ValuesQuerySet[_T, _Row]: ...
-    def complex_filter(self, filter_obj: Any) -> ValuesQuerySet[_T, _Row]: ...
-    def count(self) -> int: ...
-    def union(self, *other_qs: Any, all: bool = ...) -> ValuesQuerySet[_T, _Row]: ...
-    def intersection(self, *other_qs: Any) -> ValuesQuerySet[_T, _Row]: ...
-    def difference(self, *other_qs: Any) -> ValuesQuerySet[_T, _Row]: ...
-    def select_for_update(
-        self, nowait: bool = ..., skip_locked: bool = ..., of: Tuple = ...
-    ) -> ValuesQuerySet[_T, _Row]: ...
-    def select_related(self, *fields: Any) -> ValuesQuerySet[_T, _Row]: ...
-    def prefetch_related(self, *lookups: Any) -> ValuesQuerySet[_T, _Row]: ...
-    # TODO: return type
-    def annotate(self, *args: Any, **kwargs: Any) -> QuerySet[Any]: ...
-    def order_by(self, *field_names: Any) -> ValuesQuerySet[_T, _Row]: ...
-    def distinct(self, *field_names: Any) -> ValuesQuerySet[_T, _Row]: ...
-    # 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) -> ValuesQuerySet[_T, _Row]: ...
-    def defer(self, *fields: Any) -> ValuesQuerySet[_T, _Row]: ...
-    def only(self, *fields: Any) -> ValuesQuerySet[_T, _Row]: ...
-    def using(self, alias: Optional[str]) -> ValuesQuerySet[_T, _Row]: ...

 class RawQuerySet(Iterable[_T], Sized):
    query: RawQuery
--- a/django-stubs/forms/models.pyi
+++ b/django-stubs/forms/models.pyi
@@ -6,7 +6,7 @@ from uuid import UUID
 from django.core.files.base import File
 from django.db.models.base import Model
 from django.db.models.manager import Manager
-from django.db.models.query import QuerySet
+from django.db.models.query import QuerySet, _BaseQuerySet
 from django.db.models.query_utils import Q
 from django.forms.fields import CharField, ChoiceField, Field
 from django.forms.forms import BaseForm, DeclarativeFieldsMetaclass
@@ -262,7 +262,7 @@ class ModelMultipleChoiceField(ModelChoiceField):
    widget: Any = ...
    hidden_widget: Any = ...
    default_error_messages: Any = ...
-    def __init__(self, queryset: QuerySet, **kwargs: Any) -> None: ...
+    def __init__(self, queryset: _BaseQuerySet, **kwargs: Any) -> None: ...

 def _get_foreign_key(
    parent_model: Type[Model], model: Type[Model], fk_name: Optional[str] = ..., can_fail: bool = ...
--- a/django-stubs/views/generic/list.pyi
+++ b/django-stubs/views/generic/list.pyi
@@ -1,7 +1,7 @@
 from typing import Any, Dict, Optional, Sequence, Tuple, Type

 from django.core.paginator import Paginator
-from django.db.models.query import QuerySet
+from django.db.models.query import QuerySet, _BaseQuerySet
 from django.views.generic.base import ContextMixin, TemplateResponseMixin, View

 from django.db.models import Model
@@ -22,14 +22,14 @@ class MultipleObjectMixin(ContextMixin):
    object_list: QuerySet = ...
    def get_queryset(self) -> QuerySet: ...
    def get_ordering(self) -> Sequence[str]: ...
-    def paginate_queryset(self, queryset: QuerySet, page_size: int) -> Tuple[Paginator, int, QuerySet, bool]: ...
-    def get_paginate_by(self, queryset: QuerySet) -> Optional[int]: ...
+    def paginate_queryset(self, queryset: _BaseQuerySet, page_size: int) -> Tuple[Paginator, int, QuerySet, bool]: ...
+    def get_paginate_by(self, queryset: _BaseQuerySet) -> Optional[int]: ...
    def get_paginator(
        self, queryset: QuerySet, per_page: int, orphans: int = ..., allow_empty_first_page: bool = ..., **kwargs: Any
    ) -> Paginator: ...
    def get_paginate_orphans(self) -> int: ...
    def get_allow_empty(self) -> bool: ...
-    def get_context_object_name(self, object_list: QuerySet) -> Optional[str]: ...
+    def get_context_object_name(self, object_list: _BaseQuerySet) -> Optional[str]: ...

 class BaseListView(MultipleObjectMixin, View):
    def render_to_response(self, context: Dict[str, Any], **response_kwargs: Any) -> HttpResponse: ...