[ty] Infer single-valuedness for enums based on int/str (#19510)

## Summary

We previously didn't recognize `Literal[Color.RED]` as single-valued, if
the enum also derived from `str` or `int`:
```py
from enum import Enum

class Color(str, Enum):
    RED = "red"
    GREEN = "green"
    BLUE = "blue"

def _(color: Color):
    if color == Color.RED:
        reveal_type(color)  # previously: Color, now: Literal[Color.RED]
```

The reason for that was that `int` and `str` have "custom" `__eq__` and
`__ne__` implementations that return `bool`. We do not treat enum
literals from classes with custom `__eq__` and `__ne__` implementations
as single-valued, but of course we know that `int.__eq__` and
`str.__eq__` are well-behaved.

## Test Plan

New Markdown tests.
This commit is contained in:
David Peter 2025-07-23 15:55:42 +02:00 committed by GitHub
parent cc5885e564
commit 5a55bab3f3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 39 additions and 3 deletions

View file

@ -71,6 +71,14 @@ class CustomNeEnum(Enum):
def __ne__(self, other: object) -> bool:
return False
class StrEnum(str, Enum):
A = "a"
B = "b"
class IntEnum(int, Enum):
A = 1
B = 2
static_assert(is_single_valued(Literal[NormalEnum.NO]))
static_assert(is_single_valued(Literal[NormalEnum.YES]))
static_assert(not is_single_valued(NormalEnum))
@ -89,4 +97,12 @@ static_assert(not is_single_valued(CustomEqEnum))
static_assert(not is_single_valued(Literal[CustomNeEnum.NO]))
static_assert(not is_single_valued(Literal[CustomNeEnum.YES]))
static_assert(not is_single_valued(CustomNeEnum))
static_assert(is_single_valued(Literal[StrEnum.A]))
static_assert(is_single_valued(Literal[StrEnum.B]))
static_assert(not is_single_valued(StrEnum))
static_assert(is_single_valued(Literal[IntEnum.A]))
static_assert(is_single_valued(Literal[IntEnum.B]))
static_assert(not is_single_valued(IntEnum))
```

View file

@ -222,6 +222,9 @@ bitflags! {
///
/// This is similar to no object fallback above
const META_CLASS_NO_TYPE_FALLBACK = 1 << 2;
/// Skip looking up attributes on the builtin `int` and `str` classes.
const MRO_NO_INT_OR_STR_LOOKUP = 1 << 3;
}
}
@ -244,6 +247,11 @@ impl MemberLookupPolicy {
pub(crate) const fn meta_class_no_type_fallback(self) -> bool {
self.contains(Self::META_CLASS_NO_TYPE_FALLBACK)
}
/// Exclude attributes defined on `int` or `str` when looking up attributes.
pub(crate) const fn mro_no_int_or_str_fallback(self) -> bool {
self.contains(Self::MRO_NO_INT_OR_STR_LOOKUP)
}
}
impl Default for MemberLookupPolicy {
@ -2375,11 +2383,16 @@ impl<'db> Type<'db> {
Type::EnumLiteral(_) => {
let check_dunder = |dunder_name, allowed_return_value| {
// Note that we do explicitly exclude dunder methods on `object`, `int` and `str` here.
// The reason for this is that we know that these dunder methods behave in a predictable way.
// Only custom dunder methods need to be examined here, as they might break single-valuedness
// by always returning `False`, for example.
let call_result = self.try_call_dunder_with_policy(
db,
dunder_name,
&mut CallArguments::positional([Type::unknown()]),
MemberLookupPolicy::MRO_NO_OBJECT_FALLBACK,
MemberLookupPolicy::MRO_NO_OBJECT_FALLBACK
| MemberLookupPolicy::MRO_NO_INT_OR_STR_LOOKUP,
);
let call_result = call_result.as_ref();
call_result.is_ok_and(|bindings| {

View file

@ -1442,14 +1442,21 @@ impl<'db> ClassLiteral<'db> {
dynamic_type_to_intersect_with.get_or_insert(Type::from(superclass));
}
ClassBase::Class(class) => {
if class.is_known(db, KnownClass::Object)
let known = class.known(db);
if known == Some(KnownClass::Object)
// Only exclude `object` members if this is not an `object` class itself
&& (policy.mro_no_object_fallback() && !self.is_known(db, KnownClass::Object))
{
continue;
}
if class.is_known(db, KnownClass::Type) && policy.meta_class_no_type_fallback()
if known == Some(KnownClass::Type) && policy.meta_class_no_type_fallback() {
continue;
}
if matches!(known, Some(KnownClass::Int | KnownClass::Str))
&& policy.mro_no_int_or_str_fallback()
{
continue;
}