[ty] Expansion of enums into unions of literals

Author: sharkdp

crates/ty_python_semantic/resources/mdtest/intersection_types.md

@@ -763,6 +763,49 @@ def f(
     reveal_type(j)  # revealed: Unknown & Literal[""]
 ```
 
+## Simplification of enum literals
+
+```toml
+[environment]
+python-version = "3.12"
+```
+
+```py
+from ty_extensions import Intersection, Not
+from typing import Literal
+from enum import Enum
+
+class Color(Enum):
+    RED = "red"
+    GREEN = "green"
+    BLUE = "blue"
+
+type Red = Literal[Color.RED]
+type Green = Literal[Color.GREEN]
+type Blue = Literal[Color.BLUE]
+
+def f(
+    a: Intersection[Color, Red],
+    b: Intersection[Color, Not[Red]],
+    c: Intersection[Color, Not[Red | Green]],
+    d: Intersection[Color, Not[Red | Green | Blue]],
+    e: Intersection[Red, Not[Color]],
+    f: Intersection[Red | Green, Not[Color]],
+    g: Intersection[Not[Red], Color],
+    h: Intersection[Red, Green],
+    i: Intersection[Red | Green, Green | Blue],
+):
+    reveal_type(a)  # revealed: Literal[Color.RED]
+    reveal_type(b)  # revealed: Literal[Color.GREEN, Color.BLUE]
+    reveal_type(c)  # revealed: Literal[Color.BLUE]
+    reveal_type(d)  # revealed: Never
+    reveal_type(e)  # revealed: Never
+    reveal_type(f)  # revealed: Never
+    reveal_type(g)  # revealed: Literal[Color.GREEN, Color.BLUE]
+    reveal_type(h)  # revealed: Never
+    reveal_type(i)  # revealed: Literal[Color.GREEN]
+```
+
 ## Addition of a type to an intersection with many non-disjoint types
 
 This slightly strange-looking test is a regression test for a mistake that was nearly made in a PR:

crates/ty_python_semantic/resources/mdtest/narrow/conditionals/eq.md

@@ -35,8 +35,7 @@ class Answer(Enum):
 
 def _(answer: Answer):
     if answer != Answer.NO:
-        # TODO: This should be simplified to `Literal[Answer.YES]`
-        reveal_type(answer)  # revealed: Answer & ~Literal[Answer.NO]
+        reveal_type(answer)  # revealed: Literal[Answer.YES]
     else:
         # TODO: This should be `Literal[Answer.NO]`
         reveal_type(answer)  # revealed: Answer

crates/ty_python_semantic/resources/mdtest/narrow/conditionals/is.md

@@ -78,8 +78,7 @@ def _(answer: Answer):
     if answer is Answer.NO:
         reveal_type(answer)  # revealed: Literal[Answer.NO]
     else:
-        # TODO: This should be `Literal[Answer.YES]`
-        reveal_type(answer)  # revealed: Answer & ~Literal[Answer.NO]
+        reveal_type(answer)  # revealed: Literal[Answer.YES]
 ```
 
 ## `is` for `EllipsisType` (Python 3.10+)

crates/ty_python_semantic/resources/mdtest/union_types.md

@@ -114,6 +114,31 @@ def _(
     reveal_type(u5)  # revealed: bool | Literal[17]
 ```
 
+## Enum literals
+
+```py
+from enum import Enum
+from typing import Literal
+
+class Color(Enum):
+    RED = "red"
+    GREEN = "green"
+    BLUE = "blue"
+
+def _(
+    u1: Literal[Color.RED, Color.GREEN],
+    u2: Color | Literal[Color.RED],
+    u3: Literal[Color.RED] | Color,
+    u4: Literal[Color.RED] | Literal[Color.RED, Color.GREEN],
+    u5: Literal[Color.RED, Color.GREEN, Color.BLUE],
+) -> None:
+    reveal_type(u1)  # revealed: Literal[Color.RED, Color.GREEN]
+    reveal_type(u2)  # revealed: Color
+    reveal_type(u3)  # revealed: Color
+    reveal_type(u4)  # revealed: Literal[Color.RED, Color.GREEN]
+    reveal_type(u5)  # revealed: Color
+```
+
 ## Do not erase `Unknown`
 
 ```py

crates/ty_python_semantic/src/types.rs

@@ -8334,6 +8334,7 @@ pub struct EnumLiteralType<'db> {
     /// A reference to the enum class this literal belongs to
     enum_class: ClassLiteral<'db>,
     /// The name of the enum member
+    #[returns(ref)]
     name: Name,
 }
 

crates/ty_python_semantic/src/types/builder.rs

@@ -37,6 +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::{
     BytesLiteralType, IntersectionType, KnownClass, StringLiteralType, Type,
     TypeVarBoundOrConstraints, UnionType,
@@ -374,6 +375,42 @@ impl<'db> UnionBuilder<'db> {
                     self.elements.swap_remove(index);
                 }
             }
+            Type::EnumLiteral(enum_literal) => {
+                let enum_class = enum_literal.enum_class(self.db);
+
+                let metadata = enum_metadata(self.db, enum_class);
+                let metadata = metadata.as_ref().expect("Class of enum literal is an enum");
+
+                // TODO: This is a horribly slow implementation just to check if this works.
+                let all_members = metadata.members.iter().collect::<FxOrderSet<_>>();
+
+                let mut existing_enum_members = self
+                    .elements
+                    .iter()
+                    .filter_map(|element| {
+                        if let UnionElement::Type(Type::EnumLiteral(lit)) = element {
+                            Some(lit.name(self.db))
+                        } else {
+                            None
+                        }
+                    })
+                    .collect::<FxOrderSet<_>>();
+                existing_enum_members.insert(enum_literal.name(self.db));
+
+                if all_members == existing_enum_members {
+                    self.add_in_place(enum_literal.enum_class_instance(self.db));
+                } else {
+                    if !self.elements.iter().any(|element| match element {
+                        UnionElement::Type(ty) => {
+                            Type::EnumLiteral(enum_literal).is_subtype_of(self.db, *ty)
+                        }
+                        _ => false,
+                    }) {
+                        self.elements
+                            .push(UnionElement::Type(Type::EnumLiteral(enum_literal)));
+                    }
+                }
+            }
             // Adding `object` to a union results in `object`.
             ty if ty.is_object(self.db) => {
                 self.collapse_to_object();
@@ -773,6 +810,8 @@ impl<'db> InnerIntersectionBuilder<'db> {
                 .any(KnownClass::is_bool)
         };
 
+        let contains_enum = |enum_instance| self.positive.iter().any(|ty| *ty == enum_instance);
+
         match new_negative {
             Type::Intersection(inter) => {
                 for pos in inter.positive(db) {
@@ -814,6 +853,18 @@ impl<'db> InnerIntersectionBuilder<'db> {
             Type::AlwaysFalsy if self.positive.contains(&Type::LiteralString) => {
                 self.add_negative(db, Type::string_literal(db, ""));
             }
+            Type::EnumLiteral(enum_literal)
+                if contains_enum(enum_literal.enum_class_instance(db)) =>
+            {
+                self.add_positive(
+                    db,
+                    expand_enum_to_member_union(
+                        db,
+                        enum_literal.enum_class(db),
+                        Some(enum_literal.name(db)),
+                    ),
+                );
+            }
             _ => {
                 let mut to_remove = SmallVec::<[usize; 1]>::new();
                 for (index, existing_negative) in self.negative.iter().enumerate() {

crates/ty_python_semantic/src/types/enums.rs

@@ -5,7 +5,10 @@ use crate::{
     Db,
     place::{Place, PlaceAndQualifiers, place_from_bindings, place_from_declarations},
     semantic_index::{place_table, use_def_map},
-    types::{ClassLiteral, DynamicType, KnownClass, MemberLookupPolicy, Type, TypeQualifiers},
+    types::{
+        ClassLiteral, DynamicType, EnumLiteralType, KnownClass, MemberLookupPolicy, Type,
+        TypeQualifiers, UnionType,
+    },
 };
 
 #[derive(Debug, PartialEq, Eq, get_size2::GetSize)]
@@ -217,3 +220,22 @@ pub(crate) fn enum_metadata<'db>(
 
     Some(EnumMetadata { members, aliases })
 }
+
+pub(crate) fn expand_enum_to_member_union<'db>(
+    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()))),
+            )
+        })
+}