[ty] infer name and value for enum members (#20311)

## summary
- this pr implements the following attributes for `Enum` members:
  - `name`
  - `_name_`
  - `value`
  - `_value_`
- adds a TODO test for `my_enum_class_instance.name`
- only implements if the instance is a subclass of `Enum` re: this
[comment](https://github.com/astral-sh/ruff/pull/19481#issuecomment-3103460307)
and existing
[test](c34449ed7c/crates/ty_python_semantic/resources/mdtest/enums.md (L625))

### pointers
- https://github.com/astral-sh/ty/issues/876
- https://typing.python.org/en/latest/spec/enums.html#enum-definition
- https://github.com/astral-sh/ruff/pull/19481#issuecomment-3103460307

## test plan
- mdtests
- triaged conformance diffs here:
https://diffswarm.dev/d-01k531ag4nee3xmdeq4f3j66pb
- triaged mypy primer diffs here for django-stubs:
https://diffswarm.dev/d-01k5331n13k9yx8tvnxnkeawp3
  - added a TODO test for overriding `.value`
- discord diff seems reasonable

---------

Co-authored-by: David Peter <mail@david-peter.de>
This commit is contained in:
justin 2025-09-17 03:36:27 -04:00 committed by GitHub
parent c2fa449954
commit 9f0b942b9e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 132 additions and 28 deletions

View file

@ -4,8 +4,6 @@
)]
use std::hash::BuildHasherDefault;
use rustc_hash::FxHasher;
use crate::lint::{LintRegistry, LintRegistryBuilder};
use crate::suppression::{INVALID_IGNORE_COMMENT, UNKNOWN_RULE, UNUSED_IGNORE_COMMENT};
pub use db::Db;
@ -19,6 +17,7 @@ pub use program::{
PythonVersionWithSource, SearchPathSettings,
};
pub use python_platform::PythonPlatform;
use rustc_hash::FxHasher;
pub use semantic_model::{
Completion, CompletionKind, HasDefinition, HasType, NameKind, SemanticModel,
};

View file

@ -3528,6 +3528,29 @@ impl<'db> Type<'db> {
.value_type(db)
.member_lookup_with_policy(db, name, policy),
Type::EnumLiteral(enum_literal)
if matches!(name_str, "name" | "_name_")
&& Type::ClassLiteral(enum_literal.enum_class(db))
.is_subtype_of(db, KnownClass::Enum.to_subclass_of(db)) =>
{
Place::bound(Type::StringLiteral(StringLiteralType::new(
db,
enum_literal.name(db).as_str(),
)))
.into()
}
Type::EnumLiteral(enum_literal)
if matches!(name_str, "value" | "_value_")
&& Type::ClassLiteral(enum_literal.enum_class(db))
.is_subtype_of(db, KnownClass::Enum.to_subclass_of(db)) =>
{
enum_metadata(db, enum_literal.enum_class(db))
.and_then(|metadata| metadata.members.get(enum_literal.name(db)))
.map_or_else(|| Place::Unbound, Place::bound)
.into()
}
Type::NominalInstance(..)
| Type::ProtocolInstance(..)
| Type::BooleanLiteral(..)

View file

@ -428,9 +428,8 @@ impl<'db> UnionBuilder<'db> {
let all_members_are_in_union = metadata
.members
.difference(&enum_members_in_union)
.next()
.is_none();
.keys()
.all(|name| enum_members_in_union.contains(name));
if all_members_are_in_union {
self.add_in_place_impl(

View file

@ -724,7 +724,7 @@ impl<'db> Bindings<'db> {
db,
metadata
.members
.iter()
.keys()
.map(|member| Type::string_literal(db, member)),
)
} else {

View file

@ -2,7 +2,7 @@ use ruff_python_ast::name::Name;
use rustc_hash::FxHashMap;
use crate::{
Db, FxOrderSet,
Db, FxIndexMap,
place::{Place, PlaceAndQualifiers, place_from_bindings, place_from_declarations},
semantic_index::{place_table, use_def_map},
types::{
@ -11,24 +11,24 @@ use crate::{
},
};
#[derive(Debug, PartialEq, Eq)]
pub(crate) struct EnumMetadata {
pub(crate) members: FxOrderSet<Name>,
#[derive(Debug, PartialEq, Eq, salsa::Update)]
pub(crate) struct EnumMetadata<'db> {
pub(crate) members: FxIndexMap<Name, Type<'db>>,
pub(crate) aliases: FxHashMap<Name, Name>,
}
impl get_size2::GetSize for EnumMetadata {}
impl get_size2::GetSize for EnumMetadata<'_> {}
impl EnumMetadata {
impl EnumMetadata<'_> {
fn empty() -> Self {
EnumMetadata {
members: FxOrderSet::default(),
members: FxIndexMap::default(),
aliases: FxHashMap::default(),
}
}
pub(crate) fn resolve_member<'a>(&'a self, name: &'a Name) -> Option<&'a Name> {
if self.members.contains(name) {
if self.members.contains_key(name) {
Some(name)
} else {
self.aliases.get(name)
@ -36,18 +36,21 @@ impl EnumMetadata {
}
}
#[allow(clippy::ref_option)]
fn enum_metadata_cycle_recover(
_db: &dyn Db,
_value: &Option<EnumMetadata>,
#[allow(clippy::ref_option, clippy::trivially_copy_pass_by_ref)]
fn enum_metadata_cycle_recover<'db>(
_db: &'db dyn Db,
_value: &Option<EnumMetadata<'db>>,
_count: u32,
_class: ClassLiteral<'_>,
) -> salsa::CycleRecoveryAction<Option<EnumMetadata>> {
_class: ClassLiteral<'db>,
) -> salsa::CycleRecoveryAction<Option<EnumMetadata<'db>>> {
salsa::CycleRecoveryAction::Iterate
}
#[allow(clippy::unnecessary_wraps)]
fn enum_metadata_cycle_initial(_db: &dyn Db, _class: ClassLiteral<'_>) -> Option<EnumMetadata> {
fn enum_metadata_cycle_initial<'db>(
_db: &'db dyn Db,
_class: ClassLiteral<'db>,
) -> Option<EnumMetadata<'db>> {
Some(EnumMetadata::empty())
}
@ -57,7 +60,7 @@ fn enum_metadata_cycle_initial(_db: &dyn Db, _class: ClassLiteral<'_>) -> Option
pub(crate) fn enum_metadata<'db>(
db: &'db dyn Db,
class: ClassLiteral<'db>,
) -> Option<EnumMetadata> {
) -> Option<EnumMetadata<'db>> {
// This is a fast path to avoid traversing the MRO of known classes
if class
.known(db)
@ -217,9 +220,9 @@ pub(crate) fn enum_metadata<'db>(
}
}
Some(name.clone())
Some((name.clone(), value_ty))
})
.collect::<FxOrderSet<_>>();
.collect::<FxIndexMap<_, _>>();
if members.is_empty() {
// Enum subclasses without members are not considered enums.
@ -237,7 +240,7 @@ pub(crate) fn enum_member_literals<'a, 'db: 'a>(
enum_metadata(db, class).map(|metadata| {
metadata
.members
.iter()
.keys()
.filter(move |name| Some(*name) != exclude_member)
.map(move |name| Type::EnumLiteral(EnumLiteralType::new(db, class, name.clone())))
})