Correctly infer x: SomeEnum; x.name as union of literal names

Author: sobolevn

mypy/plugins/enums.py

@@ -54,14 +54,26 @@ def enum_name_callback(ctx: mypy.plugin.AttributeContext) -> Type:
     This plugin assumes that the provided context is an attribute access
     matching one of the strings found in 'ENUM_NAME_ACCESS'.
     """
+    # This might be `SomeEnum.Field.name` case:
     enum_field_name = _extract_underlying_field_name(ctx.type)
-    if enum_field_name is None:
-        return ctx.default_attr_type
-    else:
+    if enum_field_name is not None:
         str_type = ctx.api.named_generic_type("builtins.str", [])
         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
+    if enum_names:
+        str_type = ctx.api.named_generic_type("builtins.str", [])
+        return make_simplified_union([
+            LiteralType(enum_name, fallback=str_type)
+            for enum_name in enum_names
+        ])
+
+    return ctx.default_attr_type
+
 
 _T = TypeVar("_T")
 

test-data/unit/check-enum.test

@@ -94,6 +94,38 @@ reveal_type(Truth.true.name)    # N: Revealed type is "Literal['true']?"
 reveal_type(Truth.false.value)  # N: Revealed type is "Literal[False]?"
 [builtins fixtures/bool.pyi]
 
+[case testEnumNameValueOnType]
+from enum import Enum, Flag, IntEnum, member, nonmember
+class Colors(Enum):
+    red = 1
+    blue = 2
+    green = 3
+    white = nonmember(0)
+x: Colors = Colors.red
+reveal_type(x.name)  # N: Revealed type is "Union[Literal['red'], Literal['blue'], Literal['green']]"
+
+class Numbers(IntEnum):
+    one = 1
+    two = 2
+    three = 3
+    _magic_ = member(4)
+y: Numbers = Numbers.one
+reveal_type(y.name)  # N: Revealed type is "Union[Literal['one'], Literal['two'], Literal['three'], Literal['_magic_']]"
+
+class Flags(Flag):
+    en = "en"
+    us = "us"
+    @member
+    def all(self) -> str:
+        return "all"
+z: Flags = Flags.en
+reveal_type(z.name)  # N: Revealed type is "Union[Literal['en'], Literal['us'], Literal['all']]"
+
+class Empty(Enum): ...
+e: Empty
+reveal_type(e.name)  # N: Revealed type is "builtins.str"
+[builtins fixtures/tuple.pyi]
+
 [case testEnumValueExtended]
 from enum import Enum
 class Truth(Enum):