Fix signature of Choices member creation, add assert_type test case…

…s, run `pyright` (typeddjango#2162)

* Fix signature of Choices member creation

* Add comment regarding overloads

* Add pyright to CI, add test

* Run mypy on the new test cases

* Add more assertions, rename test folder

* Update to `pyright==1.1.364`

* Add `.gitattributes` for correct syntax highlighting

* Python compat

* [pre-commit.ci] auto fixes from pre-commit.com hooks

* type ignore comments compatibility between pyright and mypy

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

.gitattributes

@@ -0,0 +1,2 @@
+scripts/allowlist_*.txt linguist-language=ini
+pyrightconfig*.json linguist-language=jsonc

.github/workflows/test.yml

@@ -67,8 +67,10 @@ jobs:
           SETUPTOOLS_ENABLE_FEATURES=legacy-editable pip install -r ./requirements.txt
 
       # Must match `shard` definition in the test matrix:
-      - name: Run tests
+      - name: Run pytest tests
         run: PYTHONPATH='.' pytest --num-shards=4 --shard-id=${{ matrix.shard }} tests
+      - name: Run mypy on the test cases
+        run: mypy tests/assert_type
 
   stubtest:
     timeout-minutes: 10
@@ -112,12 +114,17 @@ jobs:
         run: |
           pip install -U pip setuptools wheel
           SETUPTOOLS_ENABLE_FEATURES=legacy-editable pip install -r ./requirements.txt
-      - name: Run pyright
+      - name: Run pyright on the stubs
         uses: jakebailey/pyright-action@v2
         with:
-          pylance-version: latest-release
+          version: PATH
           annotate: false
         continue-on-error: true  # TODO: remove this part
+      - name: Run pyright on the test cases
+        uses: jakebailey/pyright-action@v2
+        with:
+          version: PATH
+          project: ./pyrightconfig.testcases.json
 
   matrix-test:
     timeout-minutes: 10

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, TypeVar, overload, type_check_only
 
-from typing_extensions import Self, TypeAlias
+from _typeshed import ConvertibleToInt
+from django.utils.functional import _StrOrPromise
+from typing_extensions import TypeAlias
 
 _Self = TypeVar("_Self")
 
@@ -55,8 +57,14 @@ class _IntegerChoicesMeta(ChoicesType):
     @property
     def values(self) -> list[int]: ...
 
+# In reality, the `__init__` overloads provided below should also support
+# all the arguments of `int.__new__`/`str.__new__` (e.g. `base`, `encoding`).
+# They are omitted on purpose to avoid having convoluted stubs for these enums:
 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: ...
     @_enum_property
     def value(self) -> int: ...
 
@@ -69,6 +77,9 @@ class _TextChoicesMeta(ChoicesType):
     def values(self) -> list[str]: ...
 
 class TextChoices(Choices, StrEnum, metaclass=_TextChoicesMeta):
-    def __new__(cls, value: str | tuple[str, str]) -> Self: ...
+    @overload
+    def __init__(self, object: str) -> None: ...
+    @overload
+    def __init__(self, object: str, label: _StrOrPromise) -> None: ...
     @_enum_property
     def value(self) -> str: ...

pyproject.toml

@@ -6,28 +6,6 @@ include = '\.pyi?$'
 [tool.codespell]
 ignore-words-list = "aadd,acount,nam,asend"
 
-[tool.pyright]
-include = [
-    "django-stubs",
-    "ext/django_stubs_ext",
-    "mypy_django_plugin",
-    "scripts",
-    "tests",
-]
-exclude = [
-    ".github",
-    ".mypy_cache",
-    "build",
-]
-reportMissingTypeArgument = "warning"
-reportPrivateUsage = "none"
-stubPath = "."
-typeCheckingMode = "strict"
-
-pythonVersion = "3.8"
-pythonPlatform = "All"
-
-
 [tool.ruff]
 # Adds to default excludes: https://ruff.rs/docs/settings/#exclude
 extend-exclude = [
@@ -68,7 +46,7 @@ ignore = ["PYI021", "PYI024", "PYI041", "PYI043"]
     "F822",
     "F821",
 ]
-"tests/*.py" = ["INP001"]
+"tests/*.py" = ["INP001", "PGH003"]
 "ext/tests/*.py" = ["INP001"]
 "setup.py" = ["INP001"]
 

pyrightconfig.json

@@ -0,0 +1,23 @@
+{
+    "$schema": "https://raw.githubusercontent.com/microsoft/pyright/main/packages/vscode-pyright/schemas/pyrightconfig.schema.json",
+    "include": [
+        "django-stubs",
+        "ext/django_stubs_ext",
+        "mypy_django_plugin",
+        "scripts",
+    ],
+    "exclude": [
+        ".github",
+        ".mypy_cache",
+        "build",
+        // test cases use a custom config file
+        "tests/",
+    ],
+    "typeCheckingMode": "strict",
+    "reportMissingTypeArgument": "warning",
+    // Stubs are allowed to use private variables
+    "reportPrivateUsage": "none",
+    "stubPath": ".",
+    "pythonVersion": "3.8",
+    "pythonPlatform": "All",
+}

pyrightconfig.testcases.json

@@ -0,0 +1,23 @@
+{
+    "$schema": "https://raw.githubusercontent.com/microsoft/pyright/main/packages/vscode-pyright/schemas/pyrightconfig.schema.json",
+    "include": [
+        "tests/assert_type/"
+    ],
+    "typeCheckingMode": "strict",
+    // Extra strict settings
+    "reportShadowedImports": "error", // Don't accidentally name a file something that shadows stdlib
+    "reportImplicitStringConcatenation": "error",
+    "reportUninitializedInstanceVariable": "error",
+    "reportUnnecessaryTypeIgnoreComment": "error",
+    // Don't use '# type: ignore' to suppress with pyright
+    "enableTypeIgnoreComments": false,
+    // If a test case uses this anti-pattern, there's likely a reason and annoying to `type: ignore`.
+    // Let Ruff flag it (B006)
+    "reportCallInDefaultInitializer": "none",
+    // Too strict and not needed for type testing
+    "reportMissingSuperCall": "none",
+    // Stubs are allowed to use private variables. We may want to test those.
+    "reportPrivateUsage": "none",
+    // Stubs don't need the actual modules to be installed
+    "reportMissingModuleSource": "none",
+}

requirements.txt

@@ -15,3 +15,4 @@ Django==5.0.6; python_version >= '3.10'
 
 # Overrides:
 mypy==1.10.0
+pyright==1.1.364

tests/assert_type/db/models/check_enums.py

@@ -0,0 +1,46 @@
+from typing import List, Literal, Tuple
+
+from django.db.models import IntegerChoices, TextChoices
+from django.utils.translation import gettext_lazy as _
+from typing_extensions import assert_type
+
+
+class MyIntegerChoices(IntegerChoices):
+    A = 1
+    B = 2, "B"
+    C = 3, "B", "..."  # pyright: ignore[reportCallIssue]
+    D = 4, _("D")
+    E = 5, 1  # pyright: ignore[reportArgumentType]
+    F = "1"
+
+
+assert_type(MyIntegerChoices.A, Literal[MyIntegerChoices.A])
+assert_type(MyIntegerChoices.A.label, str)
+
+# For standard enums, type checkers may infer the type of a member's value
+# (e.g. `MyIntegerChoices.A.value` inferred as `Literal[1]`).
+# However, Django choices metaclass is using the last value for the label.
+# Type checkers relies on the stub definition of the `value` property, typed
+# as `int`/`str` for `IntegerChoices`/`TextChoices`.
+assert_type(MyIntegerChoices.A.value, int)
+
+
+class MyTextChoices(TextChoices):
+    A = "a"
+    B = "b", "B"
+    C = "c", _("C")
+    D = 1  # pyright: ignore[reportArgumentType]
+    E = "e", 1  # pyright: ignore[reportArgumentType]
+
+
+assert_type(MyTextChoices.A, Literal[MyTextChoices.A])
+assert_type(MyTextChoices.A.label, str)
+assert_type(MyTextChoices.A.value, str)
+
+
+# Assertions related to the metaclass:
+
+assert_type(MyIntegerChoices.values, List[int])
+assert_type(MyIntegerChoices.choices, List[Tuple[int, str]])
+assert_type(MyTextChoices.values, List[str])
+assert_type(MyTextChoices.choices, List[Tuple[str, str]])