Skip to content

Commit 110a1a9

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 f49a88f commit 110a1a9

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
@@ -196,22 +196,11 @@ class SomeEnum:
196196
if _implements_new(info):
197197
return ctx.default_attr_type
198198

199-
stnodes = (info.get(name) for name in info.names)
200-
201199
# Enums _can_ have methods, instance attributes, and `nonmember`s.
202200
# Omit methods and attributes created by assigning to self.*
203201
# for our value inference.
204-
node_types = (
205-
get_proper_type(n.type) if n else None
206-
for n in stnodes
207-
if n is None or not n.implicit
208-
)
209-
proper_types = [
210-
_infer_value_type_with_auto_fallback(ctx, t)
211-
for t in node_types
212-
if t is None
213-
or (not isinstance(t, CallableType) and not is_named_instance(t, "enum.nonmember"))
214-
]
202+
node_types = (get_proper_type(info[name].type) for name in info.enum_members)
203+
proper_types = [_infer_value_type_with_auto_fallback(ctx, t) for t in node_types]
215204
underlying_type = _first(proper_types)
216205
if underlying_type is None:
217206
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
@@ -2635,3 +2635,28 @@ reveal_type(Wrapper.Nested.FOO) # N: Revealed type is "Literal[__main__.Wrapper
26352635
reveal_type(Wrapper.Nested.FOO.value) # N: Revealed type is "builtins.ellipsis"
26362636
reveal_type(Wrapper.Nested.FOO._value_) # N: Revealed type is "builtins.ellipsis"
26372637
[builtins fixtures/enum.pyi]
2638+
2639+
[case testValueFallbackWithCachedMethod]
2640+
from enum import Enum, auto
2641+
from collections.abc import Hashable
2642+
from typing import Callable, Generic, TypeVar
2643+
2644+
_T = TypeVar("_T")
2645+
2646+
class _lru_cache_wrapper(Generic[_T]):
2647+
def __call__(self, *args: Hashable, **kwargs: Hashable) -> _T: ...
2648+
2649+
def cache(user_function: Callable[..., _T], /) -> _lru_cache_wrapper[_T]: ...
2650+
2651+
class Color(Enum):
2652+
RED = auto()
2653+
2654+
@cache
2655+
def lowercase_name(self) -> str:
2656+
return self.name
2657+
2658+
reveal_type(Color.RED.value) # N: Revealed type is "builtins.int"
2659+
2660+
def frobnicate(color: Color) -> None:
2661+
reveal_type(color.value) # N: Revealed type is "builtins.int"
2662+
[builtins fixtures/primitives.pyi]

0 commit comments

Comments
 (0)