From 3d9256b3c508b69426dea9d672f517459f372fa8 Mon Sep 17 00:00:00 2001
From: Nikita Sobolev <mail@sobolevn.me>
Date: Fri, 14 Jun 2024 09:10:37 +0300
Subject: [PATCH] Support `enum.nonmember` for python3.11+ (#17376)

This PR adds support for
https://docs.python.org/3.11/library/enum.html#enum.nonmember

Refs https://github.com/python/mypy/issues/12841
---
 mypy/checkmember.py              | 11 +++++++++++
 mypy/plugins/enums.py            | 15 ++++++++++++---
 test-data/unit/check-enum.test   | 28 ++++++++++++++++++++++++++++
 test-data/unit/lib-stub/enum.pyi |  5 +++++
 4 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/mypy/checkmember.py b/mypy/checkmember.py
index fa847de2e4a0..7525db25d9cd 100644
--- a/mypy/checkmember.py
+++ b/mypy/checkmember.py
@@ -1143,6 +1143,17 @@ def analyze_enum_class_attribute_access(
     if name.startswith("__") and name.replace("_", "") != "":
         return None
 
+    node = itype.type.get(name)
+    if node and node.type:
+        proper = get_proper_type(node.type)
+        # Support `A = nonmember(1)` function call and decorator.
+        if (
+            isinstance(proper, Instance)
+            and proper.type.fullname == "enum.nonmember"
+            and proper.args
+        ):
+            return proper.args[0]
+
     enum_literal = LiteralType(name, fallback=itype)
     return itype.copy_modified(last_known_value=enum_literal)
 
diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py
index 83350fe2fe11..167b330f9b09 100644
--- a/mypy/plugins/enums.py
+++ b/mypy/plugins/enums.py
@@ -20,7 +20,15 @@
 from mypy.semanal_enum import ENUM_BASES
 from mypy.subtypes import is_equivalent
 from mypy.typeops import fixup_partial_type, make_simplified_union
-from mypy.types import CallableType, Instance, LiteralType, ProperType, Type, get_proper_type
+from mypy.types import (
+    CallableType,
+    Instance,
+    LiteralType,
+    ProperType,
+    Type,
+    get_proper_type,
+    is_named_instance,
+)
 
 ENUM_NAME_ACCESS: Final = {f"{prefix}.name" for prefix in ENUM_BASES} | {
     f"{prefix}._name_" for prefix in ENUM_BASES
@@ -159,7 +167,7 @@ class SomeEnum:
 
             stnodes = (info.get(name) for name in info.names)
 
-            # Enums _can_ have methods and instance attributes.
+            # Enums _can_ have methods, instance attributes, and `nonmember`s.
             # Omit methods and attributes created by assigning to self.*
             # for our value inference.
             node_types = (
@@ -170,7 +178,8 @@ class SomeEnum:
             proper_types = [
                 _infer_value_type_with_auto_fallback(ctx, t)
                 for t in node_types
-                if t is None or not isinstance(t, CallableType)
+                if t is None
+                or (not isinstance(t, CallableType) and not is_named_instance(t, "enum.nonmember"))
             ]
             underlying_type = _first(proper_types)
             if underlying_type is None:
diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test
index e8e65f464eaf..183901416604 100644
--- a/test-data/unit/check-enum.test
+++ b/test-data/unit/check-enum.test
@@ -2138,3 +2138,31 @@ elif e == MyEnum.B:
 else:
     reveal_type(e)  # E: Statement is unreachable
 [builtins fixtures/dict.pyi]
+
+
+[case testEnumNonMemberSupport]
+# flags: --python-version 3.11
+# This was added in 3.11
+from enum import Enum, nonmember
+
+class My(Enum):
+    a = 1
+    b = 2
+    c = nonmember(3)
+
+reveal_type(My.a)  # N: Revealed type is "Literal[__main__.My.a]?"
+reveal_type(My.b)  # N: Revealed type is "Literal[__main__.My.b]?"
+reveal_type(My.c)  # N: Revealed type is "builtins.int"
+
+def accepts_my(my: My):
+    reveal_type(my.value)  # N: Revealed type is "Union[Literal[1]?, Literal[2]?]"
+
+class Other(Enum):
+    a = 1
+    @nonmember
+    class Support:
+        b = 2
+
+reveal_type(Other.a)  # N: Revealed type is "Literal[__main__.Other.a]?"
+reveal_type(Other.Support.b)  # N: Revealed type is "builtins.int"
+[builtins fixtures/dict.pyi]
diff --git a/test-data/unit/lib-stub/enum.pyi b/test-data/unit/lib-stub/enum.pyi
index 11adfc597955..32dd7c38d251 100644
--- a/test-data/unit/lib-stub/enum.pyi
+++ b/test-data/unit/lib-stub/enum.pyi
@@ -48,3 +48,8 @@ class auto(IntFlag):
 # It is python-3.11+ only:
 class StrEnum(str, Enum):
     def __new__(cls: Type[_T], value: str | _T) -> _T: ...
+
+# It is python-3.11+ only:
+class nonmember(Generic[_T]):
+    value: _T
+    def __init__(self, value: _T) -> None: ...