Skip to content

Commit

Permalink
Update models's fields' _Choices type
Browse files Browse the repository at this point in the history
Django 5 allows model field choices to be callables, mappings or
subclasses of models.Choices. This commit introduces these options in
the stubs.

Co-Authored-By: Stian Jensen <me@stianj.com>
  • Loading branch information
joaoseckler and stianjensen committed Jan 9, 2025
1 parent 93a6ef7 commit 3531bd4
Show file tree
Hide file tree
Showing 4 changed files with 110 additions and 3 deletions.
8 changes: 6 additions & 2 deletions django-stubs/db/models/fields/__init__.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,13 @@ from django import forms
from django.core import validators # due to weird mypy.stubtest error
from django.core.checks import CheckMessage
from django.db.backends.base.base import BaseDatabaseWrapper
from django.db.models import Model
from django.db.models import Choices, Model
from django.db.models.expressions import Col, Combinable, Expression, Func
from django.db.models.fields.reverse_related import ForeignObjectRel
from django.db.models.query_utils import Q, RegisterLookupMixin
from django.forms import Widget
from django.utils.choices import BlankChoiceIterator, _Choice, _ChoiceNamedGroup, _Choices, _ChoicesCallable
from django.utils.choices import BlankChoiceIterator, _Choice, _ChoiceNamedGroup, _ChoicesCallable, _ChoicesMapping
from django.utils.choices import _Choices as _ChoicesSequence
from django.utils.datastructures import DictWrapper
from django.utils.functional import _Getter, _StrOrPromise, cached_property
from typing_extensions import Self, TypeAlias
Expand All @@ -27,6 +28,9 @@ BLANK_CHOICE_DASH: list[tuple[str, str]]
_ChoicesList: TypeAlias = Sequence[_Choice] | Sequence[_ChoiceNamedGroup]
_LimitChoicesTo: TypeAlias = Q | dict[str, Any]
_LimitChoicesToCallable: TypeAlias = Callable[[], _LimitChoicesTo]
_Choices: TypeAlias = (
_ChoicesSequence | _ChoicesMapping | type[Choices] | Callable[[], _ChoicesSequence | _ChoicesMapping]
)

_F = TypeVar("_F", bound=Field, covariant=True)

Expand Down
3 changes: 2 additions & 1 deletion django-stubs/utils/choices.pyi
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
from collections.abc import Iterable, Iterator
from collections.abc import Iterable, Iterator, Mapping
from typing import Any, Protocol, TypeVar, type_check_only

from typing_extensions import TypeAlias

_Choice: TypeAlias = tuple[Any, Any]
_ChoiceNamedGroup: TypeAlias = tuple[str, Iterable[_Choice]]
_Choices: TypeAlias = Iterable[_Choice | _ChoiceNamedGroup]
_ChoicesMapping: TypeAlias = Mapping[Any, Any] # noqa: PYI047

@type_check_only
class _ChoicesCallable(Protocol):
Expand Down
89 changes: 89 additions & 0 deletions tests/assert_type/db/models/fields/test_choices.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
from collections.abc import Callable, Mapping, Sequence
from typing import TypeVar

from django.db import models
from typing_extensions import assert_type

_T = TypeVar("_T")


def to_named_seq(func: Callable[[], _T]) -> Callable[[], Sequence[tuple[str, _T]]]:
def inner() -> Sequence[tuple[str, _T]]:
return [("title", func())]

return inner


def to_named_mapping(func: Callable[[], _T]) -> Callable[[], Mapping[str, _T]]:
def inner() -> Mapping[str, _T]:
return {"title": func()}

return inner


def str_tuple() -> Sequence[tuple[str, str]]:
return (("foo", "bar"), ("fuzz", "bazz"))


def str_mapping() -> Mapping[str, str]:
return {"foo": "bar", "fuzz": "bazz"}


def int_tuple() -> Sequence[tuple[int, str]]:
return ((1, "bar"), (2, "bazz"))


def int_mapping() -> Mapping[int, str]:
return {3: "bar", 4: "bazz"}


class TestModel(models.Model):
class TextChoices(models.TextChoices):
FIRST = "foo", "bar"
SECOND = "foo2", "bar"

class IntegerChoices(models.IntegerChoices):
FIRST = 1, "bar"
SECOND = 2, "bar"

char1 = models.CharField[str, str](max_length=5, choices=TextChoices, default="foo")
char2 = models.CharField[str, str](max_length=5, choices=str_tuple, default="foo")
char3 = models.CharField[str, str](max_length=5, choices=str_mapping, default="foo")
char4 = models.CharField[str, str](max_length=5, choices=str_tuple(), default="foo")
char5 = models.CharField[str, str](max_length=5, choices=str_mapping(), default="foo")
char6 = models.CharField[str, str](max_length=5, choices=to_named_seq(str_tuple), default="foo")
char7 = models.CharField[str, str](max_length=5, choices=to_named_mapping(str_mapping), default="foo")
char8 = models.CharField[str, str](max_length=5, choices=to_named_seq(str_tuple)(), default="foo")
char9 = models.CharField[str, str](max_length=5, choices=to_named_mapping(str_mapping)(), default="foo")

int1 = models.IntegerField[int, int](choices=IntegerChoices, default=1)
int2 = models.IntegerField[int, int](choices=int_tuple, default=1)
int3 = models.IntegerField[int, int](choices=int_mapping, default=1)
int4 = models.IntegerField[int, int](choices=int_tuple(), default=1)
int5 = models.IntegerField[int, int](choices=int_mapping(), default=1)
int6 = models.IntegerField[int, int](choices=to_named_seq(int_tuple), default=1)
int7 = models.IntegerField[int, int](choices=to_named_seq(int_mapping), default=1)
int8 = models.IntegerField[int, int](choices=to_named_seq(int_tuple)(), default=1)
int9 = models.IntegerField[int, int](choices=to_named_seq(int_mapping)(), default=1)


instance = TestModel()
assert_type(instance.char1, str)
assert_type(instance.char2, str)
assert_type(instance.char3, str)
assert_type(instance.char4, str)
assert_type(instance.char5, str)
assert_type(instance.char6, str)
assert_type(instance.char7, str)
assert_type(instance.char8, str)
assert_type(instance.char9, str)

assert_type(instance.int1, int)
assert_type(instance.int2, int)
assert_type(instance.int3, int)
assert_type(instance.int4, int)
assert_type(instance.int5, int)
assert_type(instance.int6, int)
assert_type(instance.int7, int)
assert_type(instance.int8, int)
assert_type(instance.int9, int)
13 changes: 13 additions & 0 deletions tests/typecheck/db/models/test_fields.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
- case: db_models_fields_choices
main: |
from django.db import models
class MyModel(models.Model):
char1 = models.CharField[str, str](max_length=200, choices='test')
out: |
main:4: error: Argument "choices" to "CharField" has incompatible type "str"; expected "Union[Iterable[Union[Tuple[Any, Any], Tuple[str, Iterable[Tuple[Any, Any]]]]], Mapping[Any, Any], Type[Choices], Callable[[], Union[Iterable[Union[Tuple[Any, Any], Tuple[str, Iterable[Tuple[Any, Any]]]]], Mapping[Any, Any]]], None]" [arg-type]
main:4: note: Following member(s) of "str" have conflicts:
main:4: note: Expected:
main:4: note: def __iter__(self) -> Iterator[Union[Tuple[Any, Any], Tuple[str, Iterable[Tuple[Any, Any]]]]]
main:4: note: Got:
main:4: note: def __iter__(self) -> Iterator[str]

0 comments on commit 3531bd4

Please sign in to comment.