Skip to content

Commit

Permalink
Add django.db.models.query.BaseIterable and subclasses (#1242)
Browse files Browse the repository at this point in the history
Closes #1237.
  • Loading branch information
joshua-jandyco authored Nov 11, 2022
1 parent a0e98a2 commit 1b8c1b5
Show file tree
Hide file tree
Showing 2 changed files with 76 additions and 0 deletions.
27 changes: 27 additions & 0 deletions django-stubs/db/models/query.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ from typing import (
Iterator,
List,
MutableMapping,
NamedTuple,
Optional,
Reversible,
Sequence,
Expand All @@ -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]] = ...,
Expand Down
49 changes: 49 additions & 0 deletions tests/typecheck/db/models/test_query.yml
Original file line number Diff line number Diff line change
@@ -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]]")

0 comments on commit 1b8c1b5

Please sign in to comment.