Improve typing of choices (#2582)

When `__empty__` is defined, the `.choices` and `.values` properties
will add a value for the blank choice, so we expect that `None` should
be part of the return type for the value, e.g. `list[str | None]`
instead of `list[str]` for `.values` and for `.choices` we expect
`list[tuple[str | None, _StrOrPromise]` instead of
`list[tuple[str, `_StrOrPromise]]`. This is too dynamic for the stubs
and so requires some tweaks in the plugin.

* Fixed typing of `Choices` for custom mixed-in type

It's easy enough to get decent types for `IntegerChoices` and
`TextChoices` in the stubs because of subclassing `IntEnum` and
`StrEnum`, as well as having a dedicated base enum type.

But for custom choices with mixed-in base types, it's not possible to
get sensible typing as enums cannot be generic. As such, the type of the
value has to be `Any`. In the plugin, we can find the base type and use
that instead.

Author: ngnpope

django-stubs/db/models/enums.pyi

@@ -1,86 +1,83 @@
 import enum
 import sys
-from typing import Any, TypeVar, overload, type_check_only
+from typing import Any, Literal, TypeVar, overload, type_check_only
 
 from _typeshed import ConvertibleToInt
 from django.utils.functional import _StrOrPromise
 from typing_extensions import deprecated
 
-_Self = TypeVar("_Self")
-
 if sys.version_info >= (3, 11):
-    _enum_property = enum.property
-    EnumType = enum.EnumType
-    IntEnum = enum.IntEnum
-    StrEnum = enum.StrEnum
+    from enum import EnumType, IntEnum, StrEnum
+    from enum import property as enum_property
 else:
-    _enum_property = property
-    EnumType = enum.EnumMeta
+    from enum import EnumMeta as EnumType
+    from types import DynamicClassAttribute as enum_property
 
     class ReprEnum(enum.Enum): ...  # type: ignore[misc]
     class IntEnum(int, ReprEnum): ...  # type: ignore[misc]
     class StrEnum(str, ReprEnum): ...  # type: ignore[misc]
 
+_Self = TypeVar("_Self", bound=ChoicesType)
+
 class ChoicesType(EnumType):
-    # There's a contradiction between mypy and PYI019 regarding metaclasses. Where mypy
-    # disallows 'typing_extensions.Self' on metaclasses, while PYI019 try to enforce
-    # 'typing_extensions.Self' for '__new__' methods.. We've chosen to ignore the
-    # linter and trust mypy.
-    def __new__(  # noqa: PYI019
+    __empty__: _StrOrPromise
+    def __new__(
         metacls: type[_Self], classname: str, bases: tuple[type, ...], classdict: enum._EnumDict, **kwds: Any
     ) -> _Self: ...
-    def __contains__(self, member: Any) -> bool: ...
     @property
     def names(self) -> list[str]: ...
     @property
-    def choices(self) -> list[tuple[Any, str]]: ...
+    def choices(self) -> list[tuple[Any, _StrOrPromise]]: ...
     @property
-    def labels(self) -> list[str]: ...
+    def labels(self) -> list[_StrOrPromise]: ...
     @property
     def values(self) -> list[Any]: ...
+    if sys.version_info < (3, 12):
+        def __contains__(self, member: Any) -> bool: ...
 
 @deprecated("ChoicesMeta is deprecated in favor of ChoicesType and will be removed in Django 6.0.")
 class ChoicesMeta(ChoicesType): ...
 
 class Choices(enum.Enum, metaclass=ChoicesType):  # type: ignore[misc]
-    @property
-    def label(self) -> str: ...
-    @_enum_property
+    _label_: _StrOrPromise
+    do_not_call_in_templates: Literal[True]
+
+    @enum_property
+    def label(self) -> _StrOrPromise: ...
+    @enum_property
     def value(self) -> Any: ...
-    @property
-    def do_not_call_in_templates(self) -> bool: ...
 
 # fake, to keep simulate class properties
 @type_check_only
-class _IntegerChoicesMeta(ChoicesType):
+class _IntegerChoicesType(ChoicesType):
     @property
-    def choices(self) -> list[tuple[int, str]]: ...
+    def choices(self) -> list[tuple[int, _StrOrPromise]]: ...
     @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):  # type: ignore[misc]
+class IntegerChoices(Choices, IntEnum, metaclass=_IntegerChoicesType):  # type: ignore[misc]
     @overload
     def __init__(self, x: ConvertibleToInt) -> None: ...
     @overload
     def __init__(self, x: ConvertibleToInt, label: _StrOrPromise) -> None: ...
-    @_enum_property
+    @enum_property
     def value(self) -> int: ...
 
 # fake, to keep simulate class properties
 @type_check_only
-class _TextChoicesMeta(ChoicesType):
+class _TextChoicesType(ChoicesType):
     @property
-    def choices(self) -> list[tuple[str, str]]: ...
+    def choices(self) -> list[tuple[str, _StrOrPromise]]: ...
     @property
     def values(self) -> list[str]: ...
 
-class TextChoices(Choices, StrEnum, metaclass=_TextChoicesMeta):  # type: ignore[misc]
+class TextChoices(Choices, StrEnum, metaclass=_TextChoicesType):  # type: ignore[misc]
     @overload
     def __init__(self, object: str) -> None: ...
     @overload
     def __init__(self, object: str, label: _StrOrPromise) -> None: ...
-    @_enum_property
+    @enum_property
     def value(self) -> str: ...

mypy_django_plugin/lib/fullnames.py

@@ -17,6 +17,9 @@
 MANAGER_CLASS_FULLNAME = "django.db.models.manager.Manager"
 RELATED_MANAGER_CLASS = "django.db.models.fields.related_descriptors.RelatedManager"
 
+CHOICES_CLASS_FULLNAME = "django.db.models.enums.Choices"
+CHOICES_TYPE_METACLASS_FULLNAME = "django.db.models.enums.ChoicesType"
+
 WITH_ANNOTATIONS_FULLNAME = "django_stubs_ext.annotations.WithAnnotations"
 ANNOTATIONS_FULLNAME = "django_stubs_ext.annotations.Annotations"
 

mypy_django_plugin/main.py

@@ -24,6 +24,7 @@
 from mypy_django_plugin.exceptions import UnregisteredModelError
 from mypy_django_plugin.lib import fullnames, helpers
 from mypy_django_plugin.transformers import (
+    choices,
     fields,
     forms,
     init_create,
@@ -275,6 +276,12 @@ def get_attribute_hook(self, fullname: str) -> Optional[Callable[[AttributeConte
         if info and info.has_base(fullnames.STR_PROMISE_FULLNAME):
             return resolve_str_promise_attribute
 
+        if info and info.has_base(fullnames.CHOICES_TYPE_METACLASS_FULLNAME) and attr_name in ("choices", "values"):
+            return choices.transform_into_proper_attr_type
+
+        if info and info.has_base(fullnames.CHOICES_CLASS_FULLNAME) and attr_name == "value":
+            return choices.transform_into_proper_attr_type
+
         return None
 
     def get_type_analyze_hook(self, fullname: str) -> Optional[Callable[[AnalyzeTypeContext], MypyType]]:

mypy_django_plugin/transformers/choices.py

@@ -0,0 +1,87 @@
+from mypy.nodes import MemberExpr, NameExpr, TypeInfo
+from mypy.plugin import AttributeContext
+from mypy.typeanal import make_optional_type
+from mypy.types import AnyType, Instance, TupleType, TypeOfAny, get_proper_type
+from mypy.types import Type as MypyType
+
+from mypy_django_plugin.lib import fullnames, helpers
+
+
+def transform_into_proper_attr_type(ctx: AttributeContext) -> MypyType:
+    """
+    A `get_attribute_hook` to make `.choices` and `.values` optional if `__empty__` is defined.
+
+    When the `__empty__` label is specified, the `choices` and `values` properties can return a
+    blank choice/value. This hook will amend the type returned by those properties.
+    """
+    if isinstance(ctx.context, MemberExpr):
+        method_name = ctx.context.name
+    else:
+        ctx.api.fail("Unable to resolve type of choices property", ctx.context)
+        return AnyType(TypeOfAny.from_error)
+
+    expr = ctx.context.expr
+
+    if isinstance(expr, MemberExpr):
+        expr = expr.expr
+
+    if isinstance(expr, NameExpr) and isinstance(expr.node, TypeInfo):
+        node = expr.node
+    else:
+        ctx.api.fail("Unable to resolve type of choices property", ctx.context)
+        return AnyType(TypeOfAny.from_error)
+
+    default_attr_type = get_proper_type(ctx.default_attr_type)
+
+    if not node.is_enum or not node.has_base(fullnames.CHOICES_CLASS_FULLNAME):
+        return default_attr_type
+
+    # Enums with more than one base will treat the first base as the mixed-in type.
+    base_type = node.bases[0] if len(node.bases) > 1 else None
+
+    # When `__empty__` is defined, the `.choices` and `.values` properties will include `None` for
+    # the blank choice which is labelled by the value of `__empty__`.
+    has_blank_choice = node.get("__empty__") is not None
+
+    if (
+        method_name == "choices"
+        and isinstance(default_attr_type, Instance)
+        and default_attr_type.type.fullname == "builtins.list"
+        and len(default_attr_type.args) == 1
+    ):
+        choice_arg = get_proper_type(default_attr_type.args[0])
+
+        if isinstance(choice_arg, TupleType) and choice_arg.length() == 2:
+            value_arg, label_arg = choice_arg.items
+            value_arg = get_proper_type(value_arg)
+
+            if isinstance(value_arg, AnyType) and base_type is not None:
+                new_value_arg = make_optional_type(base_type) if has_blank_choice else base_type
+                new_choice_arg = choice_arg.copy_modified(items=[new_value_arg, label_arg])
+                return helpers.reparametrize_instance(default_attr_type, [new_choice_arg])
+
+            elif has_blank_choice:
+                new_value_arg = make_optional_type(value_arg)
+                new_choice_arg = choice_arg.copy_modified(items=[new_value_arg, label_arg])
+                return helpers.reparametrize_instance(default_attr_type, [new_choice_arg])
+
+    elif (
+        method_name == "values"
+        and isinstance(default_attr_type, Instance)
+        and default_attr_type.type.fullname == "builtins.list"
+        and len(default_attr_type.args) == 1
+    ):
+        value_arg = get_proper_type(default_attr_type.args[0])
+
+        if isinstance(value_arg, AnyType) and base_type is not None:
+            new_value_arg = make_optional_type(base_type) if has_blank_choice else base_type
+            return helpers.reparametrize_instance(default_attr_type, [new_value_arg])
+
+        elif has_blank_choice:
+            new_value_arg = make_optional_type(value_arg)
+            return helpers.reparametrize_instance(default_attr_type, [new_value_arg])
+
+    elif method_name == "value" and isinstance(default_attr_type, AnyType) and base_type is not None:
+        return base_type
+
+    return default_attr_type