From 50390742fb3869011b2546d1629ce3e84cd76154 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C5=81ukasz=20Langa?= Date: Fri, 8 Aug 2025 18:25:19 +0200 Subject: [PATCH] Make IntEnum/StrEnum values passable to functions expecting literal ints or strs Fixes #19616. --- mypy/subtypes.py | 24 +++++- test-data/unit/check-literal.test | 118 ++++++++++++++++++++++++++++++ 2 files changed, 141 insertions(+), 1 deletion(-) diff --git a/mypy/subtypes.py b/mypy/subtypes.py index 7da258a827f3..c1cc0947b8bd 100644 --- a/mypy/subtypes.py +++ b/mypy/subtypes.py @@ -964,7 +964,29 @@ def visit_typeddict_type(self, left: TypedDictType) -> bool: def visit_literal_type(self, left: LiteralType) -> bool: if isinstance(self.right, LiteralType): - return left == self.right + if left == self.right: + return True + # Special case: IntEnum/StrEnum literals are subtypes of int/str literals with + # the same value, e.g.: Literal[MyIntEnum.ONE] is a subtype of Literal[1] + # Literal[MyStrEnum.RED] is a subtype of Literal["red"] + # This handles IntEnum, StrEnum, and custom (int, Enum) or (str, Enum) subclasses + if ( + left.is_enum_literal() + and isinstance(left.value, str) # Enum literal values are member names + and self._is_subtype(left.fallback, self.right.fallback) + ): + # For IntEnum/StrEnum, check if the enum's actual value matches the literal + # The enum literal's value is the member name (e.g., "ONE" or "RED") + # We need to get the actual value from the enum + enum_value = left.fallback.type.get(left.value) + if enum_value is not None: + enum_type = get_proper_type(enum_value.type) + if isinstance(enum_type, Instance) and enum_type.last_known_value is not None: + # enum_type.last_known_value is the actual value for IntEnum/StrEnum + # members + if enum_type.last_known_value.value == self.right.value: + return True + return False else: return self._is_subtype(left.fallback, self.right) diff --git a/test-data/unit/check-literal.test b/test-data/unit/check-literal.test index 3c9290b8dbbb..ee074bf0c4a4 100644 --- a/test-data/unit/check-literal.test +++ b/test-data/unit/check-literal.test @@ -2995,3 +2995,121 @@ def check(obj: A[Literal[1]]) -> None: reveal_type(g('', obj)) # E: Cannot infer value of type parameter "T" of "g" \ # N: Revealed type is "Any" [builtins fixtures/tuple.pyi] + +[case testIntEnumLiteralSubtypeOfIntLiteral] +from typing import Literal +from enum import IntEnum, Enum + +class MyIntEnum(IntEnum): + ONE = 1 + TWO = 2 + THREE = 3 + +class MyCustomIntEnum(int, Enum): + ONE = 1 + TWO = 2 + THREE = 3 + +class MyEnum(Enum): + ONE = 1 + TWO = 2 + THREE = 3 + +def takes_int_literal_1(x: Literal[1]) -> None: ... +def takes_int_literal_1_2_3(x: Literal[1, 2, 3]) -> None: ... + +# IntEnum literals should be accepted where int literals are expected +takes_int_literal_1(MyIntEnum.ONE) # OK +takes_int_literal_1_2_3(MyIntEnum.TWO) # OK + +# Custom (int, Enum) literals should also be accepted +takes_int_literal_1(MyCustomIntEnum.ONE) # OK +takes_int_literal_1_2_3(MyCustomIntEnum.TWO) # OK + +# Regular Enum literals should not be accepted +takes_int_literal_1(MyEnum.ONE) # E: Argument 1 to "takes_int_literal_1" has incompatible type "Literal[MyEnum.ONE]"; expected "Literal[1]" +takes_int_literal_1_2_3(MyEnum.TWO) # E: Argument 1 to "takes_int_literal_1_2_3" has incompatible type "Literal[MyEnum.TWO]"; expected "Literal[1, 2, 3]" + +# Test assignments +x: Literal[1] = MyIntEnum.ONE # OK +y: Literal[1, 2, 3] = MyIntEnum.THREE # OK +x2: Literal[1] = MyCustomIntEnum.ONE # OK +y2: Literal[1, 2, 3] = MyCustomIntEnum.THREE # OK +z: Literal[1] = MyEnum.ONE # E: Incompatible types in assignment (expression has type "Literal[MyEnum.ONE]", variable has type "Literal[1]") + +# Test wrong values +takes_int_literal_1(MyIntEnum.TWO) # E: Argument 1 to "takes_int_literal_1" has incompatible type "Literal[MyIntEnum.TWO]"; expected "Literal[1]" +takes_int_literal_1(MyCustomIntEnum.TWO) # E: Argument 1 to "takes_int_literal_1" has incompatible type "Literal[MyCustomIntEnum.TWO]"; expected "Literal[1]" +w: Literal[1] = MyIntEnum.THREE # E: Incompatible types in assignment (expression has type "Literal[MyIntEnum.THREE]", variable has type "Literal[1]") +w2: Literal[1] = MyCustomIntEnum.THREE # E: Incompatible types in assignment (expression has type "Literal[MyCustomIntEnum.THREE]", variable has type "Literal[1]") + +# Test reverse direction - literal ints should NOT be accepted where enum is expected +def takes_int_enum(x: MyIntEnum) -> None: ... +def takes_custom_int_enum(x: MyCustomIntEnum) -> None: ... + +takes_int_enum(1) # E: Argument 1 to "takes_int_enum" has incompatible type "int"; expected "MyIntEnum" +takes_custom_int_enum(1) # E: Argument 1 to "takes_custom_int_enum" has incompatible type "int"; expected "MyCustomIntEnum" + +e1: MyIntEnum = 1 # E: Incompatible types in assignment (expression has type "int", variable has type "MyIntEnum") +e2: MyCustomIntEnum = 1 # E: Incompatible types in assignment (expression has type "int", variable has type "MyCustomIntEnum") + +[builtins fixtures/enum.pyi] + +[case testStrEnumLiteralSubtypeOfStrLiteral] +from typing import Literal +from enum import StrEnum, Enum + +class MyStrEnum(StrEnum): + RED = "red" + GREEN = "green" + BLUE = "blue" + +class MyCustomStrEnum(str, Enum): + RED = "red" + GREEN = "green" + BLUE = "blue" + +class MyEnum(Enum): + RED = "red" + GREEN = "green" + BLUE = "blue" + +def takes_str_literal_red(x: Literal["red"]) -> None: ... +def takes_str_literal_colors(x: Literal["red", "green", "blue"]) -> None: ... + +# StrEnum literals should be accepted where str literals are expected +takes_str_literal_red(MyStrEnum.RED) # OK +takes_str_literal_colors(MyStrEnum.GREEN) # OK + +# Custom (str, Enum) literals should also be accepted +takes_str_literal_red(MyCustomStrEnum.RED) # OK +takes_str_literal_colors(MyCustomStrEnum.GREEN) # OK + +# Regular Enum literals should not be accepted +takes_str_literal_red(MyEnum.RED) # E: Argument 1 to "takes_str_literal_red" has incompatible type "Literal[MyEnum.RED]"; expected "Literal['red']" +takes_str_literal_colors(MyEnum.GREEN) # E: Argument 1 to "takes_str_literal_colors" has incompatible type "Literal[MyEnum.GREEN]"; expected "Literal['red', 'green', 'blue']" + +# Test assignments +x: Literal["red"] = MyStrEnum.RED # OK +y: Literal["red", "green", "blue"] = MyStrEnum.BLUE # OK +x2: Literal["red"] = MyCustomStrEnum.RED # OK +y2: Literal["red", "green", "blue"] = MyCustomStrEnum.BLUE # OK +z: Literal["red"] = MyEnum.RED # E: Incompatible types in assignment (expression has type "Literal[MyEnum.RED]", variable has type "Literal['red']") + +# Test wrong values +takes_str_literal_red(MyStrEnum.GREEN) # E: Argument 1 to "takes_str_literal_red" has incompatible type "Literal[MyStrEnum.GREEN]"; expected "Literal['red']" +takes_str_literal_red(MyCustomStrEnum.GREEN) # E: Argument 1 to "takes_str_literal_red" has incompatible type "Literal[MyCustomStrEnum.GREEN]"; expected "Literal['red']" +w: Literal["red"] = MyStrEnum.BLUE # E: Incompatible types in assignment (expression has type "Literal[MyStrEnum.BLUE]", variable has type "Literal['red']") +w2: Literal["red"] = MyCustomStrEnum.BLUE # E: Incompatible types in assignment (expression has type "Literal[MyCustomStrEnum.BLUE]", variable has type "Literal['red']") + +# Test reverse direction - literal strings should NOT be accepted where enum is expected +def takes_str_enum(x: MyStrEnum) -> None: ... +def takes_custom_str_enum(x: MyCustomStrEnum) -> None: ... + +takes_str_enum("red") # E: Argument 1 to "takes_str_enum" has incompatible type "str"; expected "MyStrEnum" +takes_custom_str_enum("red") # E: Argument 1 to "takes_custom_str_enum" has incompatible type "str"; expected "MyCustomStrEnum" + +e1: MyStrEnum = "red" # E: Incompatible types in assignment (expression has type "str", variable has type "MyStrEnum") +e2: MyCustomStrEnum = "red" # E: Incompatible types in assignment (expression has type "str", variable has type "MyCustomStrEnum") + +[builtins fixtures/enum.pyi]