From 8ee0e48cea264c16b8e058fb1d127cbf89e1c037 Mon Sep 17 00:00:00 2001 From: Mark Gregson Date: Wed, 22 May 2024 09:28:29 +1000 Subject: [PATCH] Use an abstract QuerySet-like collection type --- django-stubs/forms/models.pyi | 4 ++-- django-stubs/utils/datastructures.pyi | 13 ++++++++++++- tests/typecheck/test_formsets.yml | 2 +- 3 files changed, 15 insertions(+), 4 deletions(-) diff --git a/django-stubs/forms/models.pyi b/django-stubs/forms/models.pyi index fab91bc9a..23b81a1d8 100644 --- a/django-stubs/forms/models.pyi +++ b/django-stubs/forms/models.pyi @@ -15,7 +15,7 @@ from django.forms.renderers import BaseRenderer from django.forms.utils import ErrorList, _DataT, _FilesT from django.forms.widgets import Widget from django.utils.choices import BaseChoiceIterator, CallableChoiceIterator, _Choices, _ChoicesCallable -from django.utils.datastructures import _PropertyDescriptor +from django.utils.datastructures import _PropertyDescriptor, _QuerySetLike from django.utils.functional import _StrOrPromise from typing_extensions import TypeAlias @@ -126,7 +126,7 @@ class BaseModelFormSet(Generic[_M, _ModelFormT], BaseFormSet[_ModelFormT]): **kwargs: Any, ) -> None: ... def initial_form_count(self) -> int: ... - def get_queryset(self) -> QuerySet[_M]: ... + def get_queryset(self) -> _QuerySetLike[_M]: ... def save_new(self, form: _ModelFormT, commit: bool = ...) -> _M: ... def save_existing(self, form: _ModelFormT, obj: _M, commit: bool = ...) -> _M: ... def delete_existing(self, obj: _M, commit: bool = ...) -> None: ... diff --git a/django-stubs/utils/datastructures.pyi b/django-stubs/utils/datastructures.pyi index 68ae74e02..f53da3844 100644 --- a/django-stubs/utils/datastructures.pyi +++ b/django-stubs/utils/datastructures.pyi @@ -1,4 +1,4 @@ -from collections.abc import Collection, Iterable, Iterator, Mapping, MutableMapping, MutableSet +from collections.abc import Collection, Iterable, Iterator, Mapping, MutableMapping, MutableSet, Sized from typing import Any, Generic, NoReturn, Protocol, TypeVar, overload, type_check_only from _typeshed import Incomplete @@ -43,6 +43,17 @@ class _IndexableCollection(Protocol[_I], Collection[_I]): # noqa: PYI046 @overload def __getitem__(self, index: slice) -> Self: ... +@type_check_only +class _QuerySetLike(Protocol[_I], Iterable[_I], Sized): # noqa: PYI046 + """ + Abstract collection type reflecting methods provided by QuerySet. + """ + + @overload + def __getitem__(self, index: int) -> _I: ... + @overload + def __getitem__(self, index: slice) -> Self: ... + class OrderedSet(MutableSet[_K]): dict: dict[_K, None] def __init__(self, iterable: Iterable[_K] | None = ...) -> None: ... diff --git a/tests/typecheck/test_formsets.yml b/tests/typecheck/test_formsets.yml index d7cc35a7c..cc93abae5 100644 --- a/tests/typecheck/test_formsets.yml +++ b/tests/typecheck/test_formsets.yml @@ -7,7 +7,7 @@ ArticleFS(instance=Article()) # E: Argument "instance" to "BaseInlineFormSet" has incompatible type "Article"; expected "Optional[Category]" [arg-type] fs = ArticleFS(instance=Category()) reveal_type(fs.instance) # N: Revealed type is "myapp.models.Category" - reveal_type(fs.get_queryset()) # N: Revealed type is "django.db.models.query.QuerySet[myapp.models.Article, myapp.models.Article]" + reveal_type(fs.get_queryset()) # N: Revealed type is "django.utils.datastructures._QuerySetLike[myapp.models.Article]" installed_apps: - myapp files: