Enum expansion in overload resolution

Author: sharkdp

crates/ty_python_semantic/resources/mdtest/call/overloads.md

@@ -400,9 +400,7 @@ def _(x: SomeEnum):
     reveal_type(f(SomeEnum.A))  # revealed: A
     reveal_type(f(SomeEnum.B))  # revealed: B
     reveal_type(f(SomeEnum.C))  # revealed: C
-    # TODO: This should not be an error. The return type should be `A | B | C` once enums are expanded
-    # error: [no-matching-overload]
-    reveal_type(f(x))  # revealed: Unknown
+    reveal_type(f(x))  # revealed: A | B | C
 ```
 
 ### No matching overloads

crates/ty_python_semantic/src/types/builder.rs

@@ -37,7 +37,7 @@
 //! are subtypes of each other (unless exactly the same literal type), we can avoid many
 //! unnecessary `is_subtype_of` checks.
 
-use crate::types::enums::{enum_metadata, expand_enum_to_member_union};
+use crate::types::enums::{enum_member_literals, enum_metadata};
 use crate::types::{
     BytesLiteralType, IntersectionType, KnownClass, StringLiteralType, Type,
     TypeVarBoundOrConstraints, UnionType,
@@ -858,10 +858,13 @@ impl<'db> InnerIntersectionBuilder<'db> {
             {
                 self.add_positive(
                     db,
-                    expand_enum_to_member_union(
+                    UnionType::from_elements(
                         db,
-                        enum_literal.enum_class(db),
-                        Some(enum_literal.name(db)),
+                        enum_member_literals(
+                            db,
+                            enum_literal.enum_class(db),
+                            Some(enum_literal.name(db)),
+                        ),
                     ),
                 );
             }

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

@@ -5,6 +5,7 @@ use ruff_python_ast as ast;
 
 use crate::Db;
 use crate::types::KnownClass;
+use crate::types::enums::{enum_member_literals, enum_metadata};
 use crate::types::tuple::{TupleSpec, TupleType};
 
 use super::Type;
@@ -199,13 +200,22 @@ impl<'a, 'db> FromIterator<(Argument<'a>, Option<Type<'db>>)> for CallArguments<
 ///
 /// Returns [`None`] if the type cannot be expanded.
 fn expand_type<'db>(db: &'db dyn Db, ty: Type<'db>) -> Option<Vec<Type<'db>>> {
-    // TODO: Expand enums to their variants
     match ty {
-        Type::NominalInstance(instance) if instance.class.is_known(db, KnownClass::Bool) => {
-            Some(vec![
-                Type::BooleanLiteral(true),
-                Type::BooleanLiteral(false),
-            ])
+        Type::NominalInstance(instance) => {
+            if instance.class.is_known(db, KnownClass::Bool) {
+                return Some(vec![
+                    Type::BooleanLiteral(true),
+                    Type::BooleanLiteral(false),
+                ]);
+            }
+
+            let class_literal = instance.class.class_literal(db).0;
+
+            if enum_metadata(db, class_literal).is_some() {
+                return Some(enum_member_literals(db, class_literal, None).collect());
+            }
+
+            None
         }
         Type::Tuple(tuple_type) => {
             // Note: This should only account for tuples of known length, i.e., `tuple[bool, ...]`

crates/ty_python_semantic/src/types/enums.rs

@@ -1,3 +1,4 @@
+use itertools::Either;
 use ruff_python_ast::name::Name;
 use rustc_hash::FxHashMap;
 
@@ -7,7 +8,7 @@ use crate::{
     semantic_index::{place_table, use_def_map},
     types::{
         ClassLiteral, DynamicType, EnumLiteralType, KnownClass, MemberLookupPolicy, Type,
-        TypeQualifiers, UnionType,
+        TypeQualifiers,
     },
 };
 
@@ -221,21 +222,20 @@ pub(crate) fn enum_metadata<'db>(
     Some(EnumMetadata { members, aliases })
 }
 
-pub(crate) fn expand_enum_to_member_union<'db>(
+pub(crate) fn enum_member_literals<'a, 'db: 'a>(
     db: &'db dyn Db,
     class: ClassLiteral<'db>,
-    exclude_member: Option<&Name>,
-) -> Type<'db> {
-    enum_metadata(db, class)
-        .as_ref()
-        .map_or(Type::Never, |metadata| {
-            UnionType::from_elements(
-                db,
-                metadata
-                    .members
-                    .iter()
-                    .filter(|name| Some(*name) != exclude_member)
-                    .map(|name| Type::EnumLiteral(EnumLiteralType::new(db, class, name.clone()))),
-            )
-        })
+    exclude_member: Option<&'a Name>,
+) -> impl Iterator<Item = Type<'a>> + 'a {
+    if let Some(metadata) = enum_metadata(db, class) {
+        Either::Left(
+            metadata
+                .members
+                .iter()
+                .filter(move |name| Some(*name) != exclude_member)
+                .map(move |name| Type::EnumLiteral(EnumLiteralType::new(db, class, name.clone()))),
+        )
+    } else {
+        Either::Right(std::iter::empty())
+    }
 }