mirror of
https://github.com/astral-sh/ruff.git
synced 2025-09-29 05:14:52 +00:00
[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:
parent
cc5885e564
commit
5a55bab3f3
3 changed files with 39 additions and 3 deletions
|
@ -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))
|
||||
```
|
||||
|
|
|
@ -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| {
|
||||
|
|
|
@ -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;
|
||||
}
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue