From bbfce88ba99e1b054e2b2d885daa0b92b1a6079d Mon Sep 17 00:00:00 2001 From: Petter Friberg Date: Mon, 6 Nov 2023 22:03:30 +0100 Subject: [PATCH] Improve types for multiple methods of `QuerySet` Also includes some methods on `Manager` as that propagates due to how they're related in the stubs/plugin. --- django-stubs/db/models/manager.pyi | 17 ++++++++++------- django-stubs/db/models/query.pyi | 17 +++++++++-------- scripts/stubtest/allowlist_todo.txt | 24 ------------------------ 3 files changed, 19 insertions(+), 39 deletions(-) diff --git a/django-stubs/db/models/manager.pyi b/django-stubs/db/models/manager.pyi index 1e5bd3eb4..f00e8acd0 100644 --- a/django-stubs/db/models/manager.pyi +++ b/django-stubs/db/models/manager.pyi @@ -4,6 +4,7 @@ from typing import Any, Generic, NoReturn, TypeVar, overload from django.db.models import Combinable from django.db.models.base import Model +from django.db.models.expressions import OrderBy from django.db.models.query import QuerySet, RawQuerySet from typing_extensions import Self @@ -18,7 +19,9 @@ class BaseManager(Generic[_T]): name: str model: type[_T] _db: str | None + def __new__(cls, *args: Any, **kwargs: Any) -> Self: ... def __init__(self) -> None: ... + def __class_getitem__(cls, *args: Any, **kwargs: Any) -> type[Self]: ... def deconstruct( self, ) -> tuple[bool, str | None, str | None, Sequence[Any] | None, dict[str, Any] | None]: ... @@ -34,7 +37,7 @@ class BaseManager(Generic[_T]): 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 iterator(self, chunk_size: int | None = ...) -> Iterator[_T]: ... def aiterator(self, chunk_size: int = ...) -> AsyncIterator[_T]: ... def aggregate(self, *args: Any, **kwargs: Any) -> dict[str, Any]: ... async def aaggregate(self, *args: Any, **kwargs: Any) -> dict[str, Any]: ... @@ -70,16 +73,16 @@ class BaseManager(Generic[_T]): async def aupdate_or_create( self, defaults: MutableMapping[str, Any] | None = ..., **kwargs: Any ) -> tuple[_T, bool]: ... - def earliest(self, *fields: Any, field_name: Any | None = ...) -> _T: ... - async def aearliest(self, *fields: Any, field_name: Any | None = ...) -> _T: ... - def latest(self, *fields: Any, field_name: Any | None = ...) -> _T: ... - async def alatest(self, *fields: Any, field_name: Any | None = ...) -> _T: ... + def earliest(self, *fields: str | OrderBy) -> _T: ... + async def aearliest(self, *fields: str | OrderBy) -> _T: ... + def latest(self, *fields: str | OrderBy) -> _T: ... + async def alatest(self, *fields: str | OrderBy) -> _T: ... def first(self) -> _T | None: ... async def afirst(self) -> _T | None: ... def last(self) -> _T | None: ... async def alast(self) -> _T | None: ... - def in_bulk(self, id_list: Iterable[Any] = ..., *, field_name: str = ...) -> dict[Any, _T]: ... - async def ain_bulk(self, id_list: Iterable[Any] = ..., *, field_name: str = ...) -> dict[Any, _T]: ... + def in_bulk(self, id_list: Iterable[Any] | None = ..., *, field_name: str = ...) -> dict[Any, _T]: ... + async def ain_bulk(self, id_list: Iterable[Any] | None = ..., *, field_name: str = ...) -> dict[Any, _T]: ... def update(self, **kwargs: Any) -> int: ... async def aupdate(self, **kwargs: Any) -> int: ... def exists(self) -> bool: ... diff --git a/django-stubs/db/models/query.pyi b/django-stubs/db/models/query.pyi index 1fe6b48e4..1202e049c 100644 --- a/django-stubs/db/models/query.pyi +++ b/django-stubs/db/models/query.pyi @@ -5,7 +5,7 @@ from typing import Any, Generic, NamedTuple, TypeVar, overload from django.db.backends.utils import _ExecuteQuery from django.db.models import Manager from django.db.models.base import Model -from django.db.models.expressions import Combinable +from django.db.models.expressions import Combinable, OrderBy from django.db.models.sql.query import Query, RawQuery from django.utils.functional import cached_property from typing_extensions import Self, TypeAlias @@ -23,6 +23,7 @@ class BaseIterable(Generic[_Row]): chunked_fetch: bool chunk_size: int def __init__(self, queryset: QuerySet[Model], chunked_fetch: bool = ..., chunk_size: int = ...) -> None: ... + def __aiter__(self) -> AsyncIterator[_Row]: ... class ModelIterable(Generic[_T], BaseIterable[_T]): def __iter__(self) -> Iterator[_T]: ... @@ -65,7 +66,7 @@ class _QuerySet(Generic[_T, _Row], Collection[_Row], Reversible[_Row], Sized): def __or__(self, other: _QuerySet[_T, _Row]) -> Self: ... # 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 iterator(self, chunk_size: int | None = ...) -> Iterator[_Row]: ... def aiterator(self, chunk_size: int = ...) -> AsyncIterator[_Row]: ... def aggregate(self, *args: Any, **kwargs: Any) -> dict[str, Any]: ... async def aaggregate(self, *args: Any, **kwargs: Any) -> dict[str, Any]: ... @@ -101,16 +102,16 @@ class _QuerySet(Generic[_T, _Row], Collection[_Row], Reversible[_Row], Sized): async def aupdate_or_create( self, defaults: MutableMapping[str, Any] | None = ..., **kwargs: Any ) -> tuple[_T, bool]: ... - def earliest(self, *fields: Any, field_name: Any | None = ...) -> _Row: ... - async def aearliest(self, *fields: Any, field_name: Any | None = ...) -> _Row: ... - def latest(self, *fields: Any, field_name: Any | None = ...) -> _Row: ... - async def alatest(self, *fields: Any, field_name: Any | None = ...) -> _Row: ... + def earliest(self, *fields: str | OrderBy) -> _Row: ... + async def aearliest(self, *fields: str | OrderBy) -> _Row: ... + def latest(self, *fields: str | OrderBy) -> _Row: ... + async def alatest(self, *fields: str | OrderBy) -> _Row: ... def first(self) -> _Row | None: ... async def afirst(self) -> _Row | None: ... def last(self) -> _Row | None: ... async def alast(self) -> _Row | None: ... - def in_bulk(self, id_list: Iterable[Any] = ..., *, field_name: str = ...) -> dict[Any, _T]: ... - async def ain_bulk(self, id_list: Iterable[Any] = ..., *, field_name: str = ...) -> dict[Any, _T]: ... + def in_bulk(self, id_list: Iterable[Any] | None = ..., *, field_name: str = ...) -> dict[Any, _T]: ... + async def ain_bulk(self, id_list: Iterable[Any] | None = ..., *, field_name: str = ...) -> dict[Any, _T]: ... def delete(self) -> tuple[int, dict[str, int]]: ... async def adelete(self) -> tuple[int, dict[str, int]]: ... def update(self, **kwargs: Any) -> int: ... diff --git a/scripts/stubtest/allowlist_todo.txt b/scripts/stubtest/allowlist_todo.txt index 1b79cbba6..cfab93bfc 100644 --- a/scripts/stubtest/allowlist_todo.txt +++ b/scripts/stubtest/allowlist_todo.txt @@ -461,14 +461,7 @@ django.contrib.gis.db.models.QuerySet.__contains__ django.contrib.gis.db.models.QuerySet.__deepcopy__ django.contrib.gis.db.models.QuerySet.__reversed__ django.contrib.gis.db.models.QuerySet.__xor__ -django.contrib.gis.db.models.QuerySet.aearliest -django.contrib.gis.db.models.QuerySet.ain_bulk -django.contrib.gis.db.models.QuerySet.alatest django.contrib.gis.db.models.QuerySet.datetimes -django.contrib.gis.db.models.QuerySet.earliest -django.contrib.gis.db.models.QuerySet.in_bulk -django.contrib.gis.db.models.QuerySet.iterator -django.contrib.gis.db.models.QuerySet.latest django.contrib.gis.db.models.RasterField.contribute_to_class django.contrib.gis.db.models.RawQuerySet django.contrib.gis.db.models.RawSQL @@ -1122,14 +1115,7 @@ django.db.models.QuerySet.__contains__ django.db.models.QuerySet.__deepcopy__ django.db.models.QuerySet.__reversed__ django.db.models.QuerySet.__xor__ -django.db.models.QuerySet.aearliest -django.db.models.QuerySet.ain_bulk -django.db.models.QuerySet.alatest django.db.models.QuerySet.datetimes -django.db.models.QuerySet.earliest -django.db.models.QuerySet.in_bulk -django.db.models.QuerySet.iterator -django.db.models.QuerySet.latest django.db.models.RawQuerySet django.db.models.RawSQL django.db.models.Ref @@ -1463,8 +1449,6 @@ django.db.models.lookups.Lookup.resolve_expression django.db.models.lookups.Lookup.select_format django.db.models.lookups.PatternLookup.process_rhs django.db.models.lookups.PostgresOperatorLookup.postgres_operator -django.db.models.manager.BaseManager.__class_getitem__ -django.db.models.manager.BaseManager.__new__ django.db.models.manager.BaseManager.aaggregate django.db.models.manager.BaseManager.abulk_create django.db.models.manager.BaseManager.abulk_update @@ -1532,20 +1516,12 @@ django.db.models.options.Options.installed django.db.models.options.Options.local_concrete_fields django.db.models.options.Options.many_to_many django.db.models.options.Options.related_objects -django.db.models.query.BaseIterable.__aiter__ django.db.models.query.EmptyQuerySet.__init__ django.db.models.query.QuerySet.__contains__ django.db.models.query.QuerySet.__deepcopy__ django.db.models.query.QuerySet.__reversed__ django.db.models.query.QuerySet.__xor__ -django.db.models.query.QuerySet.aearliest -django.db.models.query.QuerySet.ain_bulk -django.db.models.query.QuerySet.alatest django.db.models.query.QuerySet.datetimes -django.db.models.query.QuerySet.earliest -django.db.models.query.QuerySet.in_bulk -django.db.models.query.QuerySet.iterator -django.db.models.query.QuerySet.latest django.db.models.query.RawQuerySet.__aiter__ django.db.models.query.RawQuerySet.__init__ django.db.models.query.RelatedPopulator