Skip to content

Commit

Permalink
Support enum.nonmember for python3.11+ (#17376)
Browse files Browse the repository at this point in the history
  • Loading branch information
sobolevn authored Jun 14, 2024
1 parent 98a22c4 commit 3d9256b
Show file tree
Hide file tree
Showing 4 changed files with 56 additions and 3 deletions.
11 changes: 11 additions & 0 deletions mypy/checkmember.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
15 changes: 12 additions & 3 deletions mypy/plugins/enums.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 = (
Expand All @@ -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:
Expand Down
28 changes: 28 additions & 0 deletions test-data/unit/check-enum.test
Original file line number Diff line number Diff line change
Expand Up @@ -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]
5 changes: 5 additions & 0 deletions test-data/unit/lib-stub/enum.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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: ...

0 comments on commit 3d9256b

Please sign in to comment.