Use string values when checking enum TypedDict keys

mypy/checkexpr.py

@@ -4810,14 +4810,48 @@ def visit_typeddict_index_expr(
                 key_types = [typ]
 
             key_names = []
+            str_type = self.named_type("builtins.str")
             for key_type in get_proper_types(key_types):
-                if isinstance(key_type, Instance) and key_type.last_known_value is not None:
-                    key_type = key_type.last_known_value
+                seen_enum_keys: set[LiteralType] = set()
+                while True:
+                    if isinstance(key_type, Instance) and key_type.last_known_value is not None:
+                        key_type = key_type.last_known_value
+                    if (
+                        not isinstance(key_type, LiteralType)
+                        or not key_type.is_enum_literal()
+                        or not is_subtype(key_type.fallback, str_type)
+                    ):
+                        break
+                    enum_type = key_type.fallback.copy_modified(last_known_value=key_type)
+                    new_type = enum_type.type.get_containing_type_info("__new__")
+                    if key_type in seen_enum_keys or (
+                        new_type is not None
+                        and not new_type.is_enum
+                        and new_type.fullname != "builtins.str"
+                    ):
+                        self.msg.typeddict_key_must_be_string_literal(td_type, index)
+                        return AnyType(TypeOfAny.from_error), set()
+                    seen_enum_keys.add(key_type)
+                    # Enum literals store member names. Use the declared value,
+                    # without following a user-defined value property.
+                    key_type = get_proper_type(
+                        analyze_member_access(
+                            "value",
+                            enum_type,
+                            index,
+                            is_lvalue=False,
+                            is_super=False,
+                            is_operator=False,
+                            original_type=enum_type,
+                            override_info=self.chk.lookup_typeinfo("enum.Enum"),
+                            chk=self.chk,
+                        )
+                    )
 
                 if (
                     isinstance(key_type, LiteralType)
                     and isinstance(key_type.value, str)
-                    and key_type.fallback.type.fullname != "builtins.bytes"
+                    and is_subtype(key_type.fallback, str_type)
                 ):
                     key_names.append(key_type.value)
                 else:

test-data/unit/check-typeddict.test

@@ -3004,6 +3004,128 @@ d[True] # E: TypedDict key must be a string literal; expected one of ("foo")
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-typeddict.pyi]
 
+[case testTypedDictEnumKey]
+# flags: --python-version 3.11
+from enum import Enum, StrEnum
+from typing import TypedDict
+
+class PlainEnum(Enum):
+    foo = "foo"
+
+class StringEnum(StrEnum):
+    foo = "foo"
+
+class StringMixinEnum(str, Enum):
+    foo = "foo"
+
+class TD(TypedDict):
+    foo: int
+
+d: TD
+d[PlainEnum.foo]  # E: TypedDict key must be a string literal; expected one of ("foo")
+d[PlainEnum.foo] = 1  # E: TypedDict key must be a string literal; expected one of ("foo") \
+    # E: Argument 1 to "__setitem__" has incompatible type "PlainEnum"; expected "str"
+reveal_type(d[StringEnum.foo])  # N: Revealed type is "builtins.int"
+reveal_type(d[StringMixinEnum.foo])  # N: Revealed type is "builtins.int"
+[builtins fixtures/enum.pyi]
+[typing fixtures/typing-typeddict.pyi]
+
+[case testTypedDictStringEnumKeyValue]
+# flags: --python-version 3.11
+from enum import Enum, StrEnum
+from typing import Literal, TypedDict
+
+class StringEnum(StrEnum):
+    name = "value"
+    unknown = "missing"
+
+class StringMixinEnum(str, Enum):
+    name = "value"
+
+class NestedEnum(StrEnum):
+    name = StringEnum.name
+
+class TD(TypedDict):
+    name: int
+    value: str
+
+d: TD
+reveal_type(d[StringEnum.name])  # N: Revealed type is "builtins.str"
+reveal_type(d[StringMixinEnum.name])  # N: Revealed type is "builtins.str"
+reveal_type(d[NestedEnum.name])  # N: Revealed type is "builtins.str"
+d[StringEnum.name] = 1  # E: Argument 2 to "__setitem__" has incompatible type "int"; expected "str"
+d[StringEnum.unknown]  # E: TypedDict "TD" has no key "missing"
+key: Literal[StringEnum.name, StringMixinEnum.name]
+reveal_type(d[key])  # N: Revealed type is "builtins.str"
+[builtins fixtures/enum.pyi]
+[typing fixtures/typing-typeddict.pyi]
+
+[case testTypedDictStringEnumKeyValueProperty]
+# flags: --python-version 3.11
+from enum import StrEnum
+from typing import Literal, TypedDict
+
+class Key(StrEnum):
+    name = "value"
+
+    # The enum stub models value as a writable attribute.
+    @property
+    def value(self) -> Literal["name"]:  # type: ignore[override, misc]
+        return "name"
+
+class TD(TypedDict):
+    name: int
+    value: str
+
+d: TD
+reveal_type(d[Key.name])  # N: Revealed type is "builtins.str"
+reveal_type(d[Key.name.value])  # N: Revealed type is "builtins.int"
+[builtins fixtures/property.pyi]
+[typing fixtures/typing-typeddict.pyi]
+
+[case testTypedDictStringEnumKeyCustomConstructor]
+from enum import Enum
+from typing import TypedDict
+
+class StringKey(str, Enum):
+    name = "name"
+    def __new__(cls, value: str) -> "StringKey": pass
+
+class CustomString(str):
+    def __new__(cls, value: str) -> "CustomString": pass
+
+class CustomKey(CustomString, Enum):
+    name = "name"
+
+class TD(TypedDict):
+    name: int
+
+d: TD
+d[StringKey.name]  # E: TypedDict key must be a string literal; expected one of ("name")
+d[CustomKey.name]  # E: TypedDict key must be a string literal; expected one of ("name")
+[builtins fixtures/enum.pyi]
+[typing fixtures/typing-typeddict.pyi]
+
+[case testTypedDictStringEnumKeyRecursiveValue]
+# flags: --python-version 3.11
+from enum_values import Key
+from typing import TypedDict
+
+class TD(TypedDict):
+    ITEM: int
+
+d: TD
+d[Key.ITEM]  # E: TypedDict key must be a string literal; expected one of ("ITEM")
+[builtins fixtures/enum.pyi]
+[typing fixtures/typing-typeddict.pyi]
+[file enum_values.pyi]
+from enum import StrEnum
+from typing import Literal
+
+class Key(StrEnum):
+    _value_: Literal[Key.ITEM]
+    ITEM = ...
+
 [case testTypedDictUppercaseKey]
 from typing import TypedDict