Skip to content

Commit dd079fb

Browse files
kmurphy4sterliakov
andcommitted
Fix Enum.value inference for Enums with @cached methods
Before this, adding an annotation like `@functools.cache` to any method on an `Enum` caused the inference for the class's `.value` to fail (i.e., become `Any`). Fixes #19368 Co-authored-by: Stanislav Terliakov <[email protected]>
1 parent 1bf186c commit dd079fb

File tree

2 files changed

+27
-7
lines changed

2 files changed

+27
-7
lines changed

mypy/plugins/enums.py

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -184,7 +184,7 @@ class SomeEnum:
184184
if _implements_new(info):
185185
return ctx.default_attr_type
186186

187-
stnodes = (info.get(name) for name in info.names)
187+
stnodes = (info.get(name) for name in info.enum_members)
188188

189189
# Enums _can_ have methods, instance attributes, and `nonmember`s.
190190
# Omit methods and attributes created by assigning to self.*
@@ -194,12 +194,7 @@ class SomeEnum:
194194
for n in stnodes
195195
if n is None or not n.implicit
196196
)
197-
proper_types = [
198-
_infer_value_type_with_auto_fallback(ctx, t)
199-
for t in node_types
200-
if t is None
201-
or (not isinstance(t, CallableType) and not is_named_instance(t, "enum.nonmember"))
202-
]
197+
proper_types = [_infer_value_type_with_auto_fallback(ctx, t) for t in node_types]
203198
underlying_type = _first(proper_types)
204199
if underlying_type is None:
205200
return ctx.default_attr_type

test-data/unit/check-enum.test

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2539,3 +2539,28 @@ def check(thing: Things) -> None:
25392539
return None
25402540
return None # E: Statement is unreachable
25412541
[builtins fixtures/enum.pyi]
2542+
2543+
[case testValueFallbackWithCachedMethod]
2544+
from enum import Enum, auto
2545+
from collections.abc import Hashable
2546+
from typing import Callable, Generic, TypeVar
2547+
2548+
_T = TypeVar("_T")
2549+
2550+
class _lru_cache_wrapper(Generic[_T]):
2551+
def __call__(self, *args: Hashable, **kwargs: Hashable) -> _T: ...
2552+
2553+
def cache(user_function: Callable[..., _T], /) -> _lru_cache_wrapper[_T]: ...
2554+
2555+
class Color(Enum):
2556+
RED = auto()
2557+
2558+
@cache
2559+
def lowercase_name(self) -> str:
2560+
return self.name
2561+
2562+
reveal_type(Color.RED.value) # N: Revealed type is "builtins.int"
2563+
2564+
def frobnicate(color: Color) -> None:
2565+
reveal_type(color.value) # N: Revealed type is "builtins.int"
2566+
[builtins fixtures/primitives.pyi]

0 commit comments

Comments
 (0)