Skip to content

Commit e116ada

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 e116ada

File tree

2 files changed

+27
-13
lines changed

2 files changed

+27
-13
lines changed

mypy/plugins/enums.py

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

187-
stnodes = (info.get(name) for name in info.names)
188-
189187
# Enums _can_ have methods, instance attributes, and `nonmember`s.
190188
# Omit methods and attributes created by assigning to self.*
191189
# for our value inference.
192-
node_types = (
193-
get_proper_type(n.type) if n else None
194-
for n in stnodes
195-
if n is None or not n.implicit
196-
)
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-
]
190+
node_types = (get_proper_type(info[name].type) for name in info.enum_members)
191+
proper_types = [_infer_value_type_with_auto_fallback(ctx, t) for t in node_types]
203192
underlying_type = _first(proper_types)
204193
if underlying_type is None:
205194
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)