diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index c526e111243a57..4db871743f2c51 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,37 @@ reveal_type(Answer.YES.value) # revealed: Literal[1] reveal_type(Answer.NO.value) # revealed: Literal[2] ``` +Using `auto()` with non-integer mixins: + +```python +from enum import Enum, auto + +class A(str, Enum): + X = auto() + Y = auto() + +reveal_type(A.X.value) # revealed: str + +class B(bytes, Enum): + X = auto() + Y = auto() + +reveal_type(B.X.value) # revealed: bytes + +class C(tuple, Enum): + X = auto() + Y = auto() + +# TODO: this should be `tuple` +reveal_type(C.X.value) # revealed: Literal[1] + +class D(float, Enum): + X = auto() + Y = auto() + +reveal_type(D.X.value) # revealed: float +``` + Combining aliases with `auto()`: ```py diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index 7ca7c2a44a2759..15acbda09bf91b 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -4024,6 +4024,16 @@ impl<'db> Type<'db> { .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::Unbound, 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 6ddb8de18550e9..e1ce28f328d5bd 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -77,9 +77,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); @@ -150,14 +147,29 @@ pub(crate) fn enum_metadata<'db>( // enum.auto Some(KnownClass::Auto) => { auto_counter += 1; - Some(if is_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 if Type::ClassLiteral(class) + .is_subtype_of(db, KnownClass::Str.to_subclass_of(db)) + { + KnownClass::Str.to_instance(db) + } else if Type::ClassLiteral(class) + .is_subtype_of(db, KnownClass::Bytes.to_subclass_of(db)) + { + KnownClass::Bytes.to_instance(db) + } else if Type::ClassLiteral(class) + .is_subtype_of(db, KnownClass::Float.to_subclass_of(db)) + { + KnownClass::Float.to_instance(db) } else { Type::IntLiteral(auto_counter) - }) + }; + Some(auto_value_ty) } _ => None,