Support enum.member for python3.11+

mypy/plugins/default.py

@@ -41,7 +41,7 @@ class DefaultPlugin(Plugin):
     """Type checker plugin that is enabled by default."""
 
     def get_function_hook(self, fullname: str) -> Callable[[FunctionContext], Type] | None:
-        from mypy.plugins import ctypes, singledispatch
+        from mypy.plugins import ctypes, enums, singledispatch
 
         if fullname == "_ctypes.Array":
             return ctypes.array_constructor_callback
@@ -51,6 +51,8 @@ def get_function_hook(self, fullname: str) -> Callable[[FunctionContext], Type]
             import mypy.plugins.functools
 
             return mypy.plugins.functools.partial_new_callback
+        elif fullname == "enum.member":
+            return enums.enum_member_callback
 
         return None
 

mypy/plugins/enums.py

@@ -87,6 +87,8 @@ def _infer_value_type_with_auto_fallback(
         return None
     proper_type = get_proper_type(fixup_partial_type(proper_type))
     if not (isinstance(proper_type, Instance) and proper_type.type.fullname == "enum.auto"):
+        if is_named_instance(proper_type, "enum.member") and proper_type.args:
+            return proper_type.args[0]
         return proper_type
     assert isinstance(ctx.type, Instance), "An incorrect ctx.type was passed."
     info = ctx.type.type
@@ -126,6 +128,22 @@ def _implements_new(info: TypeInfo) -> bool:
     return type_with_new.fullname not in ("enum.Enum", "enum.IntEnum", "enum.StrEnum")
 
 
+def enum_member_callback(ctx: mypy.plugin.FunctionContext) -> Type:
+    """By default `member(1)` will be infered as `member[int]`,
+    we want to improve the inference to be `Literal[1]` here."""
+    if ctx.arg_types or ctx.arg_types[0]:
+        arg = get_proper_type(ctx.arg_types[0][0])
+        proper_return = get_proper_type(ctx.default_return_type)
+        if (
+            isinstance(arg, Instance)
+            and arg.last_known_value
+            and isinstance(proper_return, Instance)
+            and len(proper_return.args) == 1
+        ):
+            return proper_return.copy_modified(args=[arg])
+    return ctx.default_return_type
+
+
 def enum_value_callback(ctx: mypy.plugin.AttributeContext) -> Type:
     """This plugin refines the 'value' attribute in enums to refer to
     the original underlying value. For example, suppose we have the

test-data/unit/check-enum.test

@@ -2166,3 +2166,22 @@ class Other(Enum):
 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]
+
+
+[case testEnumMemberSupport]
+# flags: --python-version 3.11
+# This was added in 3.11
+from enum import Enum, member
+
+class A(Enum):
+    x = member(1)
+    y = 2
+
+reveal_type(A.x)  # N: Revealed type is "Literal[__main__.A.x]?"
+reveal_type(A.x.value)  # N: Revealed type is "Literal[1]?"
+reveal_type(A.y)  # N: Revealed type is "Literal[__main__.A.y]?"
+reveal_type(A.y.value)  # N: Revealed type is "Literal[2]?"
+
+def some_a(a: A):
+    reveal_type(a.value)  # N: Revealed type is "Union[Literal[1]?, Literal[2]?]"
+[builtins fixtures/dict.pyi]

test-data/unit/lib-stub/enum.pyi

@@ -53,3 +53,8 @@ class StrEnum(str, Enum):
 class nonmember(Generic[_T]):
     value: _T
     def __init__(self, value: _T) -> None: ...
+
+# It is python-3.11+ only:
+class member(Generic[_T]):
+    value: _T
+    def __init__(self, value: _T) -> None: ...