Fix signature of Choices member creation, add assert_type test cases, run pyright

[Viicos]

After seeing #2156 merged, I decided to go down the rabbit hole, as the fix is actually incorrect.

First of all, we should note that there is a discussion going on to specify the typing behavior of enums. I added a comment related to constructors, as this doesn't seem to be taken into account yet. My comment roughly describes the runtime behavior.

It might be better to wait for the chapter about enums to be merged, as we are currently using undefined behaviors. However, the stubs could easily be changed if necessary in the future.

---

Regarding the implementation, I went for an __init__ method here, as __new__ is being used to instantiate enum members when doing a lookup like Color(3).

In reality, it is also possible to do the following:

from django.db.models import IntegerChoices, TextChoices
from django.utils.translation import gettext_lazy as _

class S(TextChoices):
    A = b"bytestring", "utf-8", _("Label for A")

S.A.value
#> "bytestring"

class I(IntegerChoices):
    A  = "10", 2, _("Label for A")

I.A.value
#> 2 (converted from base 2)

TextChoices also has a different behavior from Python 3.11 onwards, as Django conditionally subclasses enum.StrEnum and enum.StrEnum has a different __new__ method (but not IntEnum for some reason). The most accurate patch would then look something like this:

diff --git a/django-stubs/db/models/enums.pyi b/django-stubs/db/models/enums.pyi
index e179273e..5745e2b5 100644
--- a/django-stubs/db/models/enums.pyi
+++ b/django-stubs/db/models/enums.pyi
@@ -1,8 +1,10 @@
 import enum
 import sys
-from typing import Any, TypeVar, type_check_only
+from typing import Any, SupportsIndex, TypeVar, overload, type_check_only
 
-from typing_extensions import Self, TypeAlias
+from _typeshed import ConvertibleToInt, ReadableBuffer
+from django.utils.functional import _StrOrPromise
+from typing_extensions import TypeAlias
 
 _Self = TypeVar("_Self")
 
@@ -56,7 +58,12 @@ class _IntegerChoicesMeta(ChoicesType):
     def values(self) -> list[int]: ...
 
 class IntegerChoices(Choices, IntEnum, metaclass=_IntegerChoicesMeta):
-    def __new__(cls, value: int) -> Self: ...
+    @overload
+    def __init__(self, x: ConvertibleToInt) -> None: ...
+    @overload
+    def __init__(self, x: ConvertibleToInt, label: _StrOrPromise) -> None: ...
+    @overload
+    def __init__(self, x: str | bytes | bytearray, base: SupportsIndex, label: _StrOrPromise) -> None: ...
     @_enum_property
     def value(self) -> int: ...
 
@@ -69,6 +76,24 @@ class _TextChoicesMeta(ChoicesType):
     def values(self) -> list[str]: ...
 
 class TextChoices(Choices, StrEnum, metaclass=_TextChoicesMeta):
-    def __new__(cls, value: str | tuple[str, str]) -> Self: ...
+    if sys.version_info >= (3, 11):
+        @overload
+        def __init__(self, object: str) -> None: ...
+        @overload
+        def __init__(self, object: str, label: _StrOrPromise) -> None: ...
+        @overload
+        def __init__(self, object: ReadableBuffer, encoding: str, label: _StrOrPromise) -> None: ...
+        @overload
+        def __init__(self, object: ReadableBuffer, encoding: str, errors: str, label: _StrOrPromise) -> None: ...
+    else:
+        @overload
+        def __init__(self, object: object) -> None: ...
+        @overload
+        def __init__(self, object: object, label: _StrOrPromise) -> None: ...
+        @overload
+        def __init__(self, object: ReadableBuffer, encoding: str, label: _StrOrPromise) -> None: ...
+        @overload
+        def __init__(self, object: ReadableBuffer, encoding: str, errors: str, label: _StrOrPromise) -> None: ...
     @_enum_property
     def value(self) -> str: ...

[sobolevn]

Please, add a test for this. Since this is quite hard to understand, really.

[Viicos]

Please, add a test for this. Since this is quite hard to understand, really.

Unfortunately not possible, this (and the previous PR) only targets pyright as mypy does not type check the arguments used when defining the enum members values (at some point -- when the spec gets updated -- it will have to).

In its current state, the PR enables the following on pyright:

from django.db.models import IntegerChoices, TextChoices
from django.utils.translation import gettext_lazy as _

class T(TextChoices):
    A = "a"  # no type checker errors
    B = "b", "B label"  # no type checker errors
    C = "c", _("C label")  # no type checker errors
    D = 1  # type checker error
    E = "e", 1  # type checker error

class I(IntegerChoices):
    A = 1  # no type checker errors
    B = "1"  # no type checker errors (everything accepted by `int(...)` works)
    C = 1, 2  # type checker error

The extra use cases I added in my original comment are really really uncommon, in fact I even wonder if someone ever used these constructs. I just added them for documentation purposes, but I don't think it's a good idea to support them.

[Viicos]

Edit: I see pyright was recently added to CI, I'll see if I can add a test anyway

[sobolevn]

Check out this script from typeshed: https://github.com/python/typeshed/blob/a9d7e861f7a46ae7acd56569326adef302e10f29/.github/workflows/tests.yml#L137-L144

I am not against adding something similar to our CI as well.

[Viicos]

Hum, so you mean writing plain Python files and make use of assert_type, as done in typeshed?

[sobolevn]

Yeap, and run the with both mypy and pyright.

[Viicos]

Sounds good, in the future I guess all tests might switch to using assert_type? Or are there still benefits to use https://pypi.org/project/pytest-mypy-plugins/?

[sobolevn]

There are still benefits in using pytest-mypy-plugins, because it offeres lots of features (like per-test mypy settings and multiple files support). But, we can migrate some tests to pure assert_type

[sobolevn]

Thanks! Great stuff!

[ngnpope]

As the author of a lot of the choices stuff I plan to review everything in the stubs related to it. Sorry I haven't got around to doing so yet.

[sobolevn]

@ngnpope thank a lot, this is very appreciated! 👍

[Viicos]

Before adding more assertions, I'll wait for a new release of pyright to land (in a week or so), so that this fix: microsoft/pyright#7941 gets included

[flaeppe]

I think this looks good. It's great that you got started on the assert_type here as after this we're able to improve the test suite!

[Viicos]

I think this looks good. It's great that you got started on the assert_type here as after this we're able to improve the test suite!

Yep, it's indeed easier to write tests (syntax highlighting, you don't depend on the phrasing of the error codes/reveal type messages). Plus it's probably way faster than the pytest plugin. I think it would be great to migrate most of the tests that don't depend on the plugin to use such assertions in the future

[flaeppe]

I'm going to go ahead and merge this, thinking that additional changes/improvements can come in any subsequent PR.

Thank you for your work done here.