Add django.db.models.query.BaseIterable and subclasses (#1242)

Closes #1237.

Author: joshua-jandyco

django-stubs/db/models/query.pyi

@@ -9,6 +9,7 @@ from typing import (
     Iterator,
     List,
     MutableMapping,
+    NamedTuple,
     Optional,
     Reversible,
     Sequence,
@@ -30,13 +31,39 @@ from django.db.models.sql.query import Query, RawQuery
 _T = TypeVar("_T", bound=Model, covariant=True)
 _Row = TypeVar("_Row", covariant=True)
 _QS = TypeVar("_QS", bound="_QuerySet")
+_TupleT = TypeVar("_TupleT", bound=tuple[Any, ...], covariant=True)
 
 MAX_GET_RESULTS: int = ...
 REPR_OUTPUT_SIZE: int = ...
 
+class BaseIterable(Generic[_Row]):
+    queryset: QuerySet[Model]
+    chunked_fetch: bool
+    chunk_size: int
+    def __init__(self, queryset: QuerySet[Model], chunked_fetch: bool = ..., chunk_size: int = ...) -> None: ...
+
+class ModelIterable(Generic[_T], BaseIterable[_T]):
+    def __iter__(self) -> Iterator[_T]: ...
+
+class RawModelIterable(BaseIterable[dict[str, Any]]):
+    def __iter__(self) -> Iterator[dict[str, Any]]: ...
+
+class ValuesIterable(BaseIterable[dict[str, Any]]):
+    def __iter__(self) -> Iterator[dict[str, Any]]: ...
+
+class ValuesListIterable(BaseIterable[_TupleT]):
+    def __iter__(self) -> Iterator[_TupleT]: ...
+
+class NamedValuesListIterable(ValuesListIterable[NamedTuple]):
+    def __iter__(self) -> Iterator[NamedTuple]: ...
+
+class FlatValuesListIterable(BaseIterable[_Row]):
+    def __iter__(self) -> Iterator[_Row]: ...
+
 class _QuerySet(Generic[_T, _Row], Collection[_Row], Reversible[_Row], Sized):
     model: Type[_T]
     query: Query
+    _iterable_class: Type[BaseIterable]
     def __init__(
         self,
         model: Optional[Type[Model]] = ...,

tests/typecheck/db/models/test_query.yml

@@ -0,0 +1,49 @@
+-   case: django_db_models_query_module_has_ModelIterable
+    main: |
+        from django.db.models.query import ModelIterable
+        from django.db.models import Model
+
+        class IntModelIterable(ModelIterable[int]):
+            pass
+
+        class NoTypeParameterModelIterable(ModelIterable):
+            pass
+
+        class ModelModelIterable(ModelIterable[Model]):
+            pass
+
+        class MyModel(Model):
+            pass
+
+        class MyModelModelIterable(ModelIterable[MyModel]):
+            pass
+    out: |
+        main:4: error: Type argument "int" of "ModelIterable" must be a subtype of "Model"
+-   case: django_db_models_query_module_has_ValuesListIterable
+    main: |
+        from django.db.models.query import ValuesListIterable
+
+        class IntValuesListIterable(ValuesListIterable[tuple[int,int]]):
+            pass
+        class StringsValuesListIterable(ValuesListIterable[tuple[str,str,str]]):
+            pass
+        class MultiTypeValuesListIterable(ValuesListIterable[tuple[str,int,float]]):
+            pass
+
+        class NonTupleValuesListIterable(ValuesListIterable[int]):
+            pass
+    out: |
+        main:10: error: Type argument "int" of "ValuesListIterable" must be a subtype of "Tuple[Any, ...]"
+-   case: django_db_models_query_module_has_NamedValuesListIterable
+    main: |
+        from django.db.models.query import NamedValuesListIterable
+    out: |
+-   case: QuerySet_has__iterable_class_defined
+    main: |
+        from django.db.models.query import QuerySet, ValuesIterable, ModelIterable
+        iterable_class = QuerySet._iterable_class
+        QuerySet._iterable_class = ValuesIterable
+        QuerySet._iterable_class = ModelIterable
+        QuerySet._iterable_class = int
+    out: |
+        main:5: error: Incompatible types in assignment (expression has type "Type[int]", variable has type "Type[BaseIterable[Any]]")