From 0320b79d43e777fbe3d978aea7233ff4264e4953 Mon Sep 17 00:00:00 2001
From: sobolevn <mail@sobolevn.me>
Date: Thu, 13 Mar 2025 12:45:24 +0300
Subject: [PATCH 1/5] Correctly infer `x: SomeEnum; x.name` as union of literal
 names

---
 mypy/plugins/enums.py          | 18 +++++++++++++++---
 test-data/unit/check-enum.test | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py
index 8b7c5df6f51f..931b6aa5fcf8 100644
--- a/mypy/plugins/enums.py
+++ b/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")
 
diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test
index a3abf53e29ac..6a0599cb2ee2 100644
--- a/test-data/unit/check-enum.test
+++ b/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):

From 5ace7b9eed2cef147189c3e581614c1e29fda6df Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 13 Mar 2025 09:48:41 +0000
Subject: [PATCH 2/5] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 mypy/plugins/enums.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py
index 931b6aa5fcf8..b8a6c107e486 100644
--- a/mypy/plugins/enums.py
+++ b/mypy/plugins/enums.py
@@ -67,10 +67,9 @@ def enum_name_callback(ctx: mypy.plugin.AttributeContext) -> 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 make_simplified_union(
+            [LiteralType(enum_name, fallback=str_type) for enum_name in enum_names]
+        )
 
     return ctx.default_attr_type
 

From efc6a0fa3f8fede35442f232337628a571eb5d9c Mon Sep 17 00:00:00 2001
From: sobolevn <mail@sobolevn.me>
Date: Thu, 13 Mar 2025 13:42:05 +0300
Subject: [PATCH 3/5] Fix tests

---
 test-data/unit/pythoneval.test | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test-data/unit/pythoneval.test b/test-data/unit/pythoneval.test
index 0e0e2b1f344d..4648afe46521 100644
--- a/test-data/unit/pythoneval.test
+++ b/test-data/unit/pythoneval.test
@@ -1710,7 +1710,7 @@ reveal_type(E.X.name)
 reveal_type(e.foo)
 reveal_type(E.Y.foo)
 [out]
-_testEnumNameWorkCorrectlyOn311.py:11: note: Revealed type is "builtins.str"
+_testEnumNameWorkCorrectlyOn311.py:11: note: Revealed type is "Union[Literal['X'], Literal['Y']]"
 _testEnumNameWorkCorrectlyOn311.py:12: note: Revealed type is "Union[Literal[1]?, Literal[2]?]"
 _testEnumNameWorkCorrectlyOn311.py:13: note: Revealed type is "Literal['X']?"
 _testEnumNameWorkCorrectlyOn311.py:14: note: Revealed type is "builtins.int"

From 97bdaffe054e1acdec09b4a9f04fef67399c4819 Mon Sep 17 00:00:00 2001
From: sobolevn <mail@sobolevn.me>
Date: Fri, 14 Mar 2025 13:56:57 +0300
Subject: [PATCH 4/5] Address review

---
 mypy/plugins/enums.py          | 33 +++++++++++++++++++++++++++++----
 test-data/unit/check-enum.test | 12 ++++++++++++
 2 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py
index b8a6c107e486..de58e5de63b7 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 @@ 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
diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test
index 6a0599cb2ee2..dc0dbf4cb536 100644
--- a/test-data/unit/check-enum.test
+++ b/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):

From 920e8a52ba42229a8185c650491856de04948d32 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 14 Mar 2025 10:58:57 +0000
Subject: [PATCH 5/5] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 mypy/plugins/enums.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py
index de58e5de63b7..269843ac69f8 100644
--- a/mypy/plugins/enums.py
+++ b/mypy/plugins/enums.py
@@ -27,10 +27,9 @@
     LiteralType,
     ProperType,
     Type,
+    UnionType,
     get_proper_type,
     is_named_instance,
-    UnionType,
-    LiteralType,
 )
 
 ENUM_NAME_ACCESS: Final = {f"{prefix}.name" for prefix in ENUM_BASES} | {
@@ -65,7 +64,9 @@ def enum_name_callback(ctx: mypy.plugin.AttributeContext) -> Type:
 
     # 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)
+    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(