mirror of
https://github.com/astral-sh/ruff.git
synced 2025-11-18 11:41:21 +00:00
[ty] Fix incorrect inference of enum.auto() for enums with non-int mixins, and imprecise inference of enum.auto() for single-member enums (#20541)
Co-authored-by: Alex Waygood <alex.waygood@gmail.com>
This commit is contained in:
parent
e4dc406a3d
commit
f63a9f2334
3 changed files with 104 additions and 7 deletions
|
|
@ -320,6 +320,11 @@ reveal_type(enum_members(Answer))
|
|||
|
||||
reveal_type(Answer.YES.value) # revealed: Literal[1]
|
||||
reveal_type(Answer.NO.value) # revealed: Literal[2]
|
||||
|
||||
class SingleMember(Enum):
|
||||
SINGLE = auto()
|
||||
|
||||
reveal_type(SingleMember.SINGLE.value) # revealed: Literal[1]
|
||||
```
|
||||
|
||||
Usages of `auto()` can be combined with manual value assignments:
|
||||
|
|
@ -348,6 +353,11 @@ class Answer(StrEnum):
|
|||
|
||||
reveal_type(Answer.YES.value) # revealed: Literal["yes"]
|
||||
reveal_type(Answer.NO.value) # revealed: Literal["no"]
|
||||
|
||||
class SingleMember(StrEnum):
|
||||
SINGLE = auto()
|
||||
|
||||
reveal_type(SingleMember.SINGLE.value) # revealed: Literal["single"]
|
||||
```
|
||||
|
||||
Using `auto()` with `IntEnum` also works as expected:
|
||||
|
|
@ -363,6 +373,52 @@ reveal_type(Answer.YES.value) # revealed: Literal[1]
|
|||
reveal_type(Answer.NO.value) # revealed: Literal[2]
|
||||
```
|
||||
|
||||
As does using `auto()` for other enums that use `int` as a mixin:
|
||||
|
||||
```py
|
||||
from enum import Enum, auto
|
||||
|
||||
class Answer(int, Enum):
|
||||
YES = auto()
|
||||
NO = auto()
|
||||
|
||||
reveal_type(Answer.YES.value) # revealed: Literal[1]
|
||||
reveal_type(Answer.NO.value) # revealed: Literal[2]
|
||||
```
|
||||
|
||||
It's [hard to predict](https://github.com/astral-sh/ruff/pull/20541#discussion_r2381878613) what the
|
||||
effect of using `auto()` will be for an arbitrary non-integer mixin, so for anything that isn't a
|
||||
`StrEnum` and has a non-`int` mixin, we simply fallback to typeshed's annotation of `Any` for the
|
||||
`value` property:
|
||||
|
||||
```python
|
||||
from enum import Enum, auto
|
||||
|
||||
class A(str, Enum):
|
||||
X = auto()
|
||||
Y = auto()
|
||||
|
||||
reveal_type(A.X.value) # revealed: Any
|
||||
|
||||
class B(bytes, Enum):
|
||||
X = auto()
|
||||
Y = auto()
|
||||
|
||||
reveal_type(B.X.value) # revealed: Any
|
||||
|
||||
class C(tuple, Enum):
|
||||
X = auto()
|
||||
Y = auto()
|
||||
|
||||
reveal_type(C.X.value) # revealed: Any
|
||||
|
||||
class D(float, Enum):
|
||||
X = auto()
|
||||
Y = auto()
|
||||
|
||||
reveal_type(D.X.value) # revealed: Any
|
||||
```
|
||||
|
||||
Combining aliases with `auto()`:
|
||||
|
||||
```py
|
||||
|
|
|
|||
|
|
@ -4365,6 +4365,16 @@ impl<'db> Type<'db> {
|
|||
Place::bound(todo_type!("ParamSpecArgs / ParamSpecKwargs")).into()
|
||||
}
|
||||
|
||||
Type::NominalInstance(instance)
|
||||
if matches!(name_str, "value" | "_value_")
|
||||
&& is_single_member_enum(db, instance.class(db).class_literal(db).0) =>
|
||||
{
|
||||
enum_metadata(db, instance.class(db).class_literal(db).0)
|
||||
.and_then(|metadata| metadata.members.get_index(0).map(|(_, v)| *v))
|
||||
.map_or(Place::Undefined, Place::bound)
|
||||
.into()
|
||||
}
|
||||
|
||||
Type::NominalInstance(..)
|
||||
| Type::ProtocolInstance(..)
|
||||
| Type::BooleanLiteral(..)
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ use crate::{
|
|||
place::{Place, PlaceAndQualifiers, place_from_bindings, place_from_declarations},
|
||||
semantic_index::{place_table, use_def_map},
|
||||
types::{
|
||||
ClassLiteral, DynamicType, EnumLiteralType, KnownClass, MemberLookupPolicy,
|
||||
ClassBase, ClassLiteral, DynamicType, EnumLiteralType, KnownClass, MemberLookupPolicy,
|
||||
StringLiteralType, Type, TypeQualifiers,
|
||||
},
|
||||
};
|
||||
|
|
@ -68,9 +68,6 @@ pub(crate) fn enum_metadata<'db>(
|
|||
return None;
|
||||
}
|
||||
|
||||
let is_str_enum =
|
||||
Type::ClassLiteral(class).is_subtype_of(db, KnownClass::StrEnum.to_subclass_of(db));
|
||||
|
||||
let scope_id = class.body_scope(db);
|
||||
let use_def_map = use_def_map(db, scope_id);
|
||||
let table = place_table(db, scope_id);
|
||||
|
|
@ -141,14 +138,48 @@ pub(crate) fn enum_metadata<'db>(
|
|||
// enum.auto
|
||||
Some(KnownClass::Auto) => {
|
||||
auto_counter += 1;
|
||||
Some(if is_str_enum {
|
||||
|
||||
// `StrEnum`s have different `auto()` behaviour to enums inheriting from `(str, Enum)`
|
||||
let auto_value_ty = if Type::ClassLiteral(class)
|
||||
.is_subtype_of(db, KnownClass::StrEnum.to_subclass_of(db))
|
||||
{
|
||||
Type::StringLiteral(StringLiteralType::new(
|
||||
db,
|
||||
name.to_lowercase().as_str(),
|
||||
))
|
||||
} else {
|
||||
Type::IntLiteral(auto_counter)
|
||||
})
|
||||
let custom_mixins: smallvec::SmallVec<[Option<KnownClass>; 1]> =
|
||||
class
|
||||
.iter_mro(db, None)
|
||||
.skip(1)
|
||||
.filter_map(ClassBase::into_class)
|
||||
.filter(|class| {
|
||||
!Type::from(*class).is_subtype_of(
|
||||
db,
|
||||
KnownClass::Enum.to_subclass_of(db),
|
||||
)
|
||||
})
|
||||
.map(|class| class.known(db))
|
||||
.filter(|class| {
|
||||
!matches!(class, Some(KnownClass::Object))
|
||||
})
|
||||
.collect();
|
||||
|
||||
// `IntEnum`s have the same `auto()` behaviour to enums inheriting from `(int, Enum)`,
|
||||
// and `IntEnum`s also have `int` in their MROs, so both cases are handled here.
|
||||
//
|
||||
// In general, the `auto()` behaviour for enums with non-`int` mixins is hard to predict,
|
||||
// so we fall back to `Any` in those cases.
|
||||
if matches!(
|
||||
custom_mixins.as_slice(),
|
||||
[] | [Some(KnownClass::Int)]
|
||||
) {
|
||||
Type::IntLiteral(auto_counter)
|
||||
} else {
|
||||
Type::any()
|
||||
}
|
||||
};
|
||||
Some(auto_value_ty)
|
||||
}
|
||||
|
||||
_ => None,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue