From f63a9f233469984408f9a855d158e5243fbd515f Mon Sep 17 00:00:00 2001 From: justin Date: Mon, 10 Nov 2025 12:53:08 -0500 Subject: [PATCH] [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 --- .../resources/mdtest/enums.md | 56 +++++++++++++++++++ crates/ty_python_semantic/src/types.rs | 10 ++++ crates/ty_python_semantic/src/types/enums.rs | 45 ++++++++++++--- 3 files changed, 104 insertions(+), 7 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index c526e11124..bdd2a92995 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -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 diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index e3326791ac..2c6f449f23 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -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(..) diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index 671b919929..dbde339221 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -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; 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,