[ty] infer type for enum members

astral-sh/ty#876

Author: thejchap

crates/ruff_memory_usage/src/lib.rs

@@ -18,3 +18,12 @@ pub fn heap_size<T: GetSize>(value: &T) -> usize {
 pub fn order_set_heap_size<T: GetSize, S>(set: &OrderSet<T, S>) -> usize {
     (set.capacity() * T::get_stack_size()) + set.iter().map(heap_size).sum::<usize>()
 }
+
+/// An implementation of [`GetSize::get_heap_size`] for [`OrderMap`].
+pub fn order_map_heap_size<K: GetSize, V: GetSize, S>(map: &ordermap::OrderMap<K, V, S>) -> usize {
+    (map.capacity() * (K::get_stack_size() + V::get_stack_size()))
+        + map
+            .iter()
+            .map(|(k, v)| heap_size(k) + heap_size(v))
+            .sum::<usize>()
+}

crates/ty_python_semantic/resources/mdtest/enums.md

@@ -4,6 +4,7 @@
 
 ```py
 from enum import Enum
+from typing import Literal
 
 class Color(Enum):
     RED = 1
@@ -13,6 +14,12 @@ class Color(Enum):
 reveal_type(Color.RED)  # revealed: Literal[Color.RED]
 # TODO: This could be `Literal[1]`
 reveal_type(Color.RED.value)  # revealed: Any
+# TODO: This could be `Literal[1]`
+reveal_type(Color.RED._value_)  # revealed: Any
+
+reveal_type(Color.RED.name)  # revealed: Literal["RED"]
+
+reveal_type(Color.RED._name_)  # revealed: Literal["RED"]
 
 # TODO: Should be `Color` or `Literal[Color.RED]`
 reveal_type(Color["RED"])  # revealed: Unknown
@@ -21,6 +28,13 @@ reveal_type(Color["RED"])  # revealed: Unknown
 reveal_type(Color(1))  # revealed: Color
 
 reveal_type(Color.RED in Color)  # revealed: bool
+
+def func1(red_or_blue: Literal[Color.RED, Color.BLUE]):
+    reveal_type(red_or_blue.name)  # revealed: Literal["RED", "BLUE"]
+
+def func2(any_color: Color):
+    # TODO: Literal["RED", "GREEN", "BLUE"]
+    reveal_type(any_color.name)  # revealed: Any
 ```
 
 ## Enum members

crates/ty_python_semantic/src/types.rs

@@ -3465,6 +3465,14 @@ impl<'db> Type<'db> {
                 .value_type(db)
                 .member_lookup_with_policy(db, name, policy),
 
+            Type::EnumLiteral(enum_literal) if matches!(name_str, "name" | "_name_") => {
+                Place::bound(Type::StringLiteral(StringLiteralType::new(
+                    db,
+                    enum_literal.name(db).as_str(),
+                )))
+                .into()
+            }
+
             Type::NominalInstance(..)
             | Type::ProtocolInstance(..)
             | Type::BooleanLiteral(..)
@@ -3582,7 +3590,7 @@ impl<'db> Type<'db> {
                     _ => None,
                 } {
                     if let Some(metadata) = enum_metadata(db, enum_class) {
-                        if let Some(resolved_name) = metadata.resolve_member(&name) {
+                        if let Some(resolved_name) = metadata.resolve_member(db, &name) {
                             return Place::Type(
                                 Type::EnumLiteral(EnumLiteralType::new(
                                     db,

crates/ty_python_semantic/src/types/builder.rs

@@ -428,7 +428,10 @@ impl<'db> UnionBuilder<'db> {
                     .collect::<FxOrderSet<_>>();
 
                 let all_members_are_in_union = metadata
-                    .members
+                    .members(self.db)
+                    .keys()
+                    .cloned()
+                    .collect::<FxOrderSet<_>>()
                     .difference(&enum_members_in_union)
                     .next()
                     .is_none();

crates/ty_python_semantic/src/types/call/bind.rs

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

crates/ty_python_semantic/src/types/enums.rs

@@ -1,8 +1,10 @@
+use std::collections::BTreeMap;
+
 use ruff_python_ast::name::Name;
 use rustc_hash::FxHashMap;
 
 use crate::{
-    Db, FxOrderSet,
+    Db, FxOrderMap,
     place::{Place, PlaceAndQualifiers, place_from_bindings, place_from_declarations},
     semantic_index::{place_table, use_def_map},
     types::{
@@ -11,44 +13,52 @@ use crate::{
     },
 };
 
-#[derive(Debug, PartialEq, Eq)]
-pub(crate) struct EnumMetadata {
-    pub(crate) members: FxOrderSet<Name>,
-    pub(crate) aliases: FxHashMap<Name, Name>,
+#[salsa::interned(debug, heap_size=EnumMetadata::heap_size)]
+pub(crate) struct EnumMetadata<'db> {
+    #[returns(ref)]
+    pub(crate) members: FxOrderMap<Name, Type<'db>>,
+    #[returns(ref)]
+    pub(crate) aliases: BTreeMap<Name, Name>,
 }
 
-impl get_size2::GetSize for EnumMetadata {}
+impl get_size2::GetSize for EnumMetadata<'_> {}
 
-impl EnumMetadata {
-    fn empty() -> Self {
-        EnumMetadata {
-            members: FxOrderSet::default(),
-            aliases: FxHashMap::default(),
-        }
+impl<'db> EnumMetadata<'db> {
+    fn empty(db: &'db dyn Db) -> Self {
+        EnumMetadata::new(db, FxOrderMap::default(), BTreeMap::default())
     }
 
-    pub(crate) fn resolve_member<'a>(&'a self, name: &'a Name) -> Option<&'a Name> {
-        if self.members.contains(name) {
+    pub(crate) fn resolve_member<'a>(&'a self, db: &'a dyn Db, name: &'a Name) -> Option<&'a Name> {
+        if self.members(db).contains_key(name) {
             Some(name)
         } else {
-            self.aliases.get(name)
+            self.aliases(db).get(name)
         }
     }
+
+    fn heap_size(
+        (members, aliases): &(FxOrderMap<Name, Type<'db>>, BTreeMap<Name, Name>),
+    ) -> usize {
+        ruff_memory_usage::order_map_heap_size(members) + ruff_memory_usage::heap_size(aliases)
+    }
 }
 
-#[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> {
-    Some(EnumMetadata::empty())
+fn enum_metadata_cycle_initial<'db>(
+    db: &'db dyn Db,
+    _class: ClassLiteral<'db>,
+) -> Option<EnumMetadata<'db>> {
+    Some(EnumMetadata::empty(db))
 }
 
 /// List all members of an enum.
@@ -57,7 +67,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)
@@ -97,7 +107,7 @@ pub(crate) fn enum_metadata<'db>(
         None
     };
 
-    let mut aliases = FxHashMap::default();
+    let mut aliases = BTreeMap::default();
 
     let members = use_def_map
         .all_end_of_scope_symbol_bindings()
@@ -217,16 +227,16 @@ pub(crate) fn enum_metadata<'db>(
                 }
             }
 
-            Some(name.clone())
+            Some((name.clone(), value_ty))
         })
-        .collect::<FxOrderSet<_>>();
+        .collect::<FxOrderMap<_, _>>();
 
     if members.is_empty() {
         // Enum subclasses without members are not considered enums.
         return None;
     }
 
-    Some(EnumMetadata { members, aliases })
+    Some(EnumMetadata::new(db, members, aliases))
 }
 
 pub(crate) fn enum_member_literals<'a, 'db: 'a>(
@@ -236,15 +246,15 @@ pub(crate) fn enum_member_literals<'a, 'db: 'a>(
 ) -> Option<impl Iterator<Item = Type<'a>> + 'a> {
     enum_metadata(db, class).map(|metadata| {
         metadata
-            .members
-            .iter()
+            .members(db)
+            .keys()
             .filter(move |name| Some(*name) != exclude_member)
             .map(move |name| Type::EnumLiteral(EnumLiteralType::new(db, class, name.clone())))
     })
 }
 
 pub(crate) fn is_single_member_enum<'db>(db: &'db dyn Db, class: ClassLiteral<'db>) -> bool {
-    enum_metadata(db, class).is_some_and(|metadata| metadata.members.len() == 1)
+    enum_metadata(db, class).is_some_and(|metadata| metadata.members(db).len() == 1)
 }
 
 pub(crate) fn is_enum_class<'db>(db: &'db dyn Db, ty: Type<'db>) -> bool {