diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index 2c3c03375c..c526e11124 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -265,6 +265,41 @@ reveal_type(enum_members(Color)) reveal_type(Color.red) ``` +Multiple aliases to the same member are also supported. This is a regression test for +: + +```py +from ty_extensions import enum_members + +class ManyAliases(Enum): + real_member = "real_member" + alias1 = "real_member" + alias2 = "real_member" + alias3 = "real_member" + + other_member = "other_real_member" + +# revealed: tuple[Literal["real_member"], Literal["other_member"]] +reveal_type(enum_members(ManyAliases)) + +reveal_type(ManyAliases.real_member) # revealed: Literal[ManyAliases.real_member] +reveal_type(ManyAliases.alias1) # revealed: Literal[ManyAliases.real_member] +reveal_type(ManyAliases.alias2) # revealed: Literal[ManyAliases.real_member] +reveal_type(ManyAliases.alias3) # revealed: Literal[ManyAliases.real_member] + +reveal_type(ManyAliases.real_member.value) # revealed: Literal["real_member"] +reveal_type(ManyAliases.real_member.name) # revealed: Literal["real_member"] + +reveal_type(ManyAliases.alias1.value) # revealed: Literal["real_member"] +reveal_type(ManyAliases.alias1.name) # revealed: Literal["real_member"] + +reveal_type(ManyAliases.alias2.value) # revealed: Literal["real_member"] +reveal_type(ManyAliases.alias2.name) # revealed: Literal["real_member"] + +reveal_type(ManyAliases.alias3.value) # revealed: Literal["real_member"] +reveal_type(ManyAliases.alias3.name) # revealed: Literal["real_member"] +``` + ### Using `auto()` ```toml diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index 47000ce7ea..6ddb8de185 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -194,12 +194,19 @@ pub(crate) fn enum_metadata<'db>( // we don't know if it's a duplicate or not. if matches!( value_ty, - Type::IntLiteral(_) | Type::StringLiteral(_) | Type::BytesLiteral(_) + Type::BooleanLiteral(_) + | Type::IntLiteral(_) + | Type::StringLiteral(_) + | Type::BytesLiteral(_) ) { - if let Some(previous) = enum_values.insert(value_ty, name.clone()) { - aliases.insert(name.clone(), previous); + if let Some(canonical) = enum_values.get(&value_ty) { + // This is a duplicate value, create an alias to the canonical (first) member + aliases.insert(name.clone(), canonical.clone()); return None; } + + // This is the first occurrence of this value, track it as the canonical member + enum_values.insert(value_ty, name.clone()); } let declarations = use_def_map.end_of_scope_symbol_declarations(symbol_id);