Address review

Author: sobolevn

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 @@ def enum_name_callback(ctx: mypy.plugin.AttributeContext) -> Type:
         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 @@ def _extract_underlying_field_name(typ: Type) -> str | None:
     # 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

test-data/unit/check-enum.test

@@ -126,6 +126,18 @@ e: Empty
 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):