mirror of
https://github.com/astral-sh/ruff.git
synced 2025-09-29 13:24:57 +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:
|
def __ne__(self, other: object) -> bool:
|
||||||
return False
|
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.NO]))
|
||||||
static_assert(is_single_valued(Literal[NormalEnum.YES]))
|
static_assert(is_single_valued(Literal[NormalEnum.YES]))
|
||||||
static_assert(not is_single_valued(NormalEnum))
|
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.NO]))
|
||||||
static_assert(not is_single_valued(Literal[CustomNeEnum.YES]))
|
static_assert(not is_single_valued(Literal[CustomNeEnum.YES]))
|
||||||
static_assert(not is_single_valued(CustomNeEnum))
|
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
|
/// This is similar to no object fallback above
|
||||||
const META_CLASS_NO_TYPE_FALLBACK = 1 << 2;
|
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 {
|
pub(crate) const fn meta_class_no_type_fallback(self) -> bool {
|
||||||
self.contains(Self::META_CLASS_NO_TYPE_FALLBACK)
|
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 {
|
impl Default for MemberLookupPolicy {
|
||||||
|
@ -2375,11 +2383,16 @@ impl<'db> Type<'db> {
|
||||||
|
|
||||||
Type::EnumLiteral(_) => {
|
Type::EnumLiteral(_) => {
|
||||||
let check_dunder = |dunder_name, allowed_return_value| {
|
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(
|
let call_result = self.try_call_dunder_with_policy(
|
||||||
db,
|
db,
|
||||||
dunder_name,
|
dunder_name,
|
||||||
&mut CallArguments::positional([Type::unknown()]),
|
&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();
|
let call_result = call_result.as_ref();
|
||||||
call_result.is_ok_and(|bindings| {
|
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));
|
dynamic_type_to_intersect_with.get_or_insert(Type::from(superclass));
|
||||||
}
|
}
|
||||||
ClassBase::Class(class) => {
|
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
|
// Only exclude `object` members if this is not an `object` class itself
|
||||||
&& (policy.mro_no_object_fallback() && !self.is_known(db, KnownClass::Object))
|
&& (policy.mro_no_object_fallback() && !self.is_known(db, KnownClass::Object))
|
||||||
{
|
{
|
||||||
continue;
|
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;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue