[ty] List all enum members (#19283)

## Summary

Adds a way to list all members of an `Enum` and implements almost all of
the mechanisms by which members are distinguished from non-members
([spec](https://typing.python.org/en/latest/spec/enums.html#defining-members)).
This has no effect on actual enums, so far.

## Test Plan

New Markdown tests using `ty_extensions.enum_members`.
This commit is contained in:
David Peter
2025-07-14 13:18:17 +02:00
committed by GitHub
parent cb530a0216
commit f22da352db
7 changed files with 639 additions and 2 deletions

View File

@@ -69,6 +69,7 @@ mod context;
mod cyclic;
mod diagnostic;
mod display;
mod enums;
mod function;
mod generics;
pub(crate) mod ide_support;

View File

@@ -31,7 +31,7 @@ use crate::types::tuple::TupleType;
use crate::types::{
BoundMethodType, ClassLiteral, DataclassParams, KnownClass, KnownInstanceType,
MethodWrapperKind, PropertyInstanceType, SpecialFormType, TypeMapping, UnionType,
WrapperDescriptorKind, ide_support, todo_type,
WrapperDescriptorKind, enums, ide_support, todo_type,
};
use ruff_db::diagnostic::{Annotation, Diagnostic, Severity, SubDiagnostic};
use ruff_python_ast as ast;
@@ -661,6 +661,22 @@ impl<'db> Bindings<'db> {
}
}
Some(KnownFunction::EnumMembers) => {
if let [Some(ty)] = overload.parameter_types() {
let return_ty = match ty {
Type::ClassLiteral(class) => TupleType::from_elements(
db,
enums::enum_members(db, *class)
.into_iter()
.map(|member| Type::string_literal(db, &member)),
),
_ => Type::unknown(),
};
overload.set_return_type(return_ty);
}
}
Some(KnownFunction::AllMembers) => {
if let [Some(ty)] = overload.parameter_types() {
overload.set_return_type(TupleType::from_elements(

View File

@@ -2365,6 +2365,9 @@ pub enum KnownClass {
Super,
// enum
Enum,
Auto,
Member,
Nonmember,
// abc
ABCMeta,
// Types
@@ -2485,6 +2488,9 @@ impl KnownClass {
| Self::Deque
| Self::Float
| Self::Enum
| Self::Auto
| Self::Member
| Self::Nonmember
| Self::ABCMeta
| Self::Iterable
// Empty tuples are AlwaysFalse; non-empty tuples are AlwaysTrue
@@ -2563,6 +2569,9 @@ impl KnownClass {
Self::ABCMeta
| Self::Any
| Self::Enum
| Self::Auto
| Self::Member
| Self::Nonmember
| Self::ChainMap
| Self::Exception
| Self::ExceptionGroup
@@ -2643,6 +2652,9 @@ impl KnownClass {
| Self::Deque
| Self::OrderedDict
| Self::Enum
| Self::Auto
| Self::Member
| Self::Nonmember
| Self::ABCMeta
| Self::Super
| Self::StdlibAlias
@@ -2708,6 +2720,9 @@ impl KnownClass {
Self::Deque => "deque",
Self::OrderedDict => "OrderedDict",
Self::Enum => "Enum",
Self::Auto => "auto",
Self::Member => "member",
Self::Nonmember => "nonmember",
Self::ABCMeta => "ABCMeta",
Self::Super => "super",
Self::Iterable => "Iterable",
@@ -2929,7 +2944,7 @@ impl KnownClass {
| Self::Property => KnownModule::Builtins,
Self::VersionInfo => KnownModule::Sys,
Self::ABCMeta => KnownModule::Abc,
Self::Enum => KnownModule::Enum,
Self::Enum | Self::Auto | Self::Member | Self::Nonmember => KnownModule::Enum,
Self::GenericAlias
| Self::ModuleType
| Self::FunctionType
@@ -3042,6 +3057,9 @@ impl KnownClass {
| Self::ParamSpecKwargs
| Self::TypeVarTuple
| Self::Enum
| Self::Auto
| Self::Member
| Self::Nonmember
| Self::ABCMeta
| Self::Super
| Self::NamedTuple
@@ -3110,6 +3128,9 @@ impl KnownClass {
| Self::ParamSpecKwargs
| Self::TypeVarTuple
| Self::Enum
| Self::Auto
| Self::Member
| Self::Nonmember
| Self::ABCMeta
| Self::Super
| Self::UnionType
@@ -3182,6 +3203,9 @@ impl KnownClass {
"_NoDefaultType" => Self::NoDefaultType,
"SupportsIndex" => Self::SupportsIndex,
"Enum" => Self::Enum,
"auto" => Self::Auto,
"member" => Self::Member,
"nonmember" => Self::Nonmember,
"ABCMeta" => Self::ABCMeta,
"super" => Self::Super,
"_version_info" => Self::VersionInfo,
@@ -3243,6 +3267,9 @@ impl KnownClass {
| Self::MethodType
| Self::MethodWrapperType
| Self::Enum
| Self::Auto
| Self::Member
| Self::Nonmember
| Self::ABCMeta
| Self::Super
| Self::NotImplementedType
@@ -3762,6 +3789,7 @@ mod tests {
KnownClass::BaseExceptionGroup | KnownClass::ExceptionGroup => PythonVersion::PY311,
KnownClass::GenericAlias => PythonVersion::PY39,
KnownClass::KwOnly => PythonVersion::PY310,
KnownClass::Member | KnownClass::Nonmember => PythonVersion::PY311,
_ => PythonVersion::PY37,
};

View File

@@ -0,0 +1,145 @@
use rustc_hash::FxHashSet;
use crate::{
Db,
place::{Place, place_from_bindings, place_from_declarations},
semantic_index::{place_table, use_def_map},
types::{ClassLiteral, KnownClass, MemberLookupPolicy, Type},
};
/// List all members of an enum.
pub(crate) fn enum_members<'db>(db: &'db dyn Db, class: ClassLiteral<'db>) -> Vec<String> {
let scope_id = class.body_scope(db);
let use_def_map = use_def_map(db, scope_id);
let table = place_table(db, scope_id);
let mut enum_values: FxHashSet<Type<'db>> = FxHashSet::default();
// TODO: handle `StrEnum` which uses lowercase names as values when using `auto()`.
let mut auto_counter = 0;
let ignored_names: Option<Vec<&str>> = if let Some(ignore) = table.place_id_by_name("_ignore_")
{
let ignore_bindings = use_def_map.all_reachable_bindings(ignore);
let ignore_place = place_from_bindings(db, ignore_bindings);
match ignore_place {
Place::Type(Type::StringLiteral(ignored_names), _) => {
Some(ignored_names.value(db).split_ascii_whitespace().collect())
}
// TODO: support the list-variant of `_ignore_`.
_ => None,
}
} else {
None
};
use_def_map
.all_end_of_scope_bindings()
.filter_map(|(place_id, bindings)| {
let name = table
.place_expr(place_id)
.as_name()
.map(ToString::to_string)?;
if name.starts_with("__") && !name.ends_with("__") {
// Skip private attributes
return None;
}
if name == "_ignore_"
|| ignored_names
.as_ref()
.is_some_and(|names| names.contains(&name.as_str()))
{
// Skip ignored attributes
return None;
}
let inferred = place_from_bindings(db, bindings);
let value_ty = match inferred {
Place::Unbound => {
return None;
}
Place::Type(ty, _) => {
match ty {
Type::Callable(_) | Type::FunctionLiteral(_) => {
// Some types are specifically disallowed for enum members.
return None;
}
// enum.nonmember
Type::NominalInstance(instance)
if instance.class.is_known(db, KnownClass::Nonmember) =>
{
return None;
}
// enum.member
Type::NominalInstance(instance)
if instance.class.is_known(db, KnownClass::Member) =>
{
ty.member(db, "value")
.place
.ignore_possibly_unbound()
.unwrap_or(Type::unknown())
}
// enum.auto
Type::NominalInstance(instance)
if instance.class.is_known(db, KnownClass::Auto) =>
{
auto_counter += 1;
Type::IntLiteral(auto_counter)
}
_ => {
let dunder_get = ty
.member_lookup_with_policy(
db,
"__get__".into(),
MemberLookupPolicy::NO_INSTANCE_FALLBACK,
)
.place;
match dunder_get {
Place::Unbound | Place::Type(Type::Dynamic(_), _) => ty,
Place::Type(_, _) => {
// Descriptors are not considered members.
return None;
}
}
}
}
}
};
// Duplicate values are aliases that are not considered separate members. This check is only
// performed if we can infer a precise literal type for the enum member. If we only get `int`,
// we don't know if it's a duplicate or not.
if matches!(
value_ty,
Type::IntLiteral(_) | Type::StringLiteral(_) | Type::BytesLiteral(_)
) && !enum_values.insert(value_ty)
{
return None;
}
let declarations = use_def_map.end_of_scope_declarations(place_id);
let declared = place_from_declarations(db, declarations);
match declared.map(|d| d.place) {
Ok(Place::Unbound) => {
// Undeclared attributes are considered members
}
Ok(Place::Type(Type::NominalInstance(instance), _))
if instance.class.is_known(db, KnownClass::Member) =>
{
// If the attribute is specifically declared with `enum.member`, it is considered a member
}
_ => {
// Declared attributes are considered non-members
return None;
}
}
Some(name)
})
.collect()
}

View File

@@ -957,6 +957,8 @@ pub enum KnownFunction {
GenericContext,
/// `ty_extensions.dunder_all_names`
DunderAllNames,
/// `ty_extensions.enum_members`
EnumMembers,
/// `ty_extensions.all_members`
AllMembers,
/// `ty_extensions.top_materialization`
@@ -1025,6 +1027,7 @@ impl KnownFunction {
| Self::BottomMaterialization
| Self::GenericContext
| Self::DunderAllNames
| Self::EnumMembers
| Self::StaticAssert
| Self::AllMembers => module.is_ty_extensions(),
Self::ImportModule => module.is_importlib(),
@@ -1288,6 +1291,7 @@ pub(crate) mod tests {
| KnownFunction::IsSubtypeOf
| KnownFunction::GenericContext
| KnownFunction::DunderAllNames
| KnownFunction::EnumMembers
| KnownFunction::StaticAssert
| KnownFunction::IsDisjointFrom
| KnownFunction::IsSingleValued