Address review
diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index b8a6c10..de58e5d 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py
@@ -29,6 +29,8 @@ Type, get_proper_type, is_named_instance, + UnionType, + LiteralType, ) ENUM_NAME_ACCESS: Final = {f"{prefix}.name" for prefix in ENUM_BASES} | { @@ -61,10 +63,9 @@ literal_type = LiteralType(enum_field_name, fallback=str_type) return str_type.copy_modified(last_known_value=literal_type) - # Or `field: SomeEnum = SomeEnum.field; field.name` case: - if not isinstance(ctx.type, Instance) or not ctx.type.type.is_enum: - return ctx.default_attr_type - enum_names = ctx.type.type.enum_members + # Or `field: SomeEnum = SomeEnum.field; field.name` case, + # Or `field: Literal[Some.A, Some.B]; field.name` case: + enum_names = _extract_enum_names_from_type(ctx.type) or _extract_enum_names_from_literal_union(ctx.type) if enum_names: str_type = ctx.api.named_generic_type("builtins.str", []) return make_simplified_union( @@ -296,3 +297,27 @@ # as a string. assert isinstance(underlying_literal.value, str) return underlying_literal.value + + +def _extract_enum_names_from_type(typ: ProperType) -> list[str] | None: + if not isinstance(typ, Instance) or not typ.type.is_enum: + return None + return typ.type.enum_members + + +def _extract_enum_names_from_literal_union(typ: ProperType) -> list[str] | None: + if not isinstance(typ, UnionType): + return None + + names = [] + for item in typ.relevant_items(): + pitem = get_proper_type(item) + if isinstance(pitem, Instance) and pitem.last_known_value and pitem.type.is_enum: + assert isinstance(pitem.last_known_value.value, str) + names.append(pitem.last_known_value.value) + elif isinstance(pitem, LiteralType): + assert isinstance(pitem.value, str) + names.append(pitem.value) + else: + return None + return names
diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test index 6a0599c..dc0dbf4 100644 --- a/test-data/unit/check-enum.test +++ b/test-data/unit/check-enum.test
@@ -126,6 +126,18 @@ reveal_type(e.name) # N: Revealed type is "builtins.str" [builtins fixtures/tuple.pyi] +[case testEnumNameValueOnUnionOfLiteral] +from enum import Enum +from typing import Literal +class Colors(Enum): + red = 1 + blue = 2 + green = 3 + +color: Literal[Colors.red, Colors.blue] +reveal_type(color.name) # N: Revealed type is "Union[Literal['red'], Literal['blue']]" +[builtins fixtures/tuple.pyi] + [case testEnumValueExtended] from enum import Enum class Truth(Enum):