[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:
justin 2025-11-10 12:53:08 -05:00 committed by GitHub
parent e4dc406a3d
commit f63a9f2334
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 104 additions and 7 deletions

View file

@ -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

View file

@ -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(..)

View file

@ -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,