Skip to content

[ty] Simplify unions of enum literals and subtypes thereof #20324

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Sep 10, 2025
Merged

[ty] Simplify unions of enum literals and subtypes thereof #20324

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion crates/ty_python_semantic/resources/mdtest/union_types.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,7 +118,8 @@ def _(

```py
from enum import Enum
from typing import Literal
from typing import Literal, Any
from ty_extensions import Intersection

class Color(Enum):
RED = "red"
Expand All @@ -139,6 +140,13 @@ def _(
reveal_type(u4) # revealed: Literal[Color.RED, Color.GREEN]
reveal_type(u5) # revealed: Color
reveal_type(u6) # revealed: Color

def _(
u1: Intersection[Literal[Color.RED], Any] | Literal[Color.RED],
u2: Literal[Color.RED] | Intersection[Literal[Color.RED], Any],
):
reveal_type(u1) # revealed: Literal[Color.RED]
reveal_type(u2) # revealed: Literal[Color.RED]
```

## Do not erase `Unknown`
Expand Down
143 changes: 73 additions & 70 deletions crates/ty_python_semantic/src/types/builder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -444,93 +444,96 @@ impl<'db> UnionBuilder<'db> {
.filter_map(UnionElement::to_type_element)
.any(|ty| Type::EnumLiteral(enum_member_to_add).is_subtype_of(self.db, ty))
{
self.elements
.push(UnionElement::Type(Type::EnumLiteral(enum_member_to_add)));
Comment on lines -447 to -448
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was the problem. Instead of just using self.elements.push, we should go through the whole simplification logic in the _ => {…} branch, which I have now moved to a separate function. Every other change in this file is just the "moving out to a separate function" part.

self.push_type(Type::EnumLiteral(enum_member_to_add), seen_aliases);
}
}
// Adding `object` to a union results in `object`.
ty if ty.is_object(self.db) => {
self.collapse_to_object();
}
_ => {
let bool_pair = if let Type::BooleanLiteral(b) = ty {
Some(Type::BooleanLiteral(!b))
} else {
None
};
self.push_type(ty, seen_aliases);
}
}
}

// If an alias gets here, it means we aren't unpacking aliases, and we also
// shouldn't try to simplify aliases out of the union, because that will require
// unpacking them.
let should_simplify_full = !matches!(ty, Type::TypeAlias(_));
fn push_type(&mut self, ty: Type<'db>, seen_aliases: &mut Vec<Type<'db>>) {
let bool_pair = if let Type::BooleanLiteral(b) = ty {
Some(Type::BooleanLiteral(!b))
} else {
None
};

let mut to_remove = SmallVec::<[usize; 2]>::new();
let ty_negated = if should_simplify_full {
ty.negate(self.db)
} else {
Type::Never // won't be used
};
// If an alias gets here, it means we aren't unpacking aliases, and we also
// shouldn't try to simplify aliases out of the union, because that will require
// unpacking them.
let should_simplify_full = !matches!(ty, Type::TypeAlias(_));

for (index, element) in self.elements.iter_mut().enumerate() {
let element_type = match element.try_reduce(self.db, ty) {
ReduceResult::KeepIf(keep) => {
if !keep {
to_remove.push(index);
}
continue;
}
ReduceResult::Type(ty) => ty,
ReduceResult::CollapseToObject => {
self.collapse_to_object();
return;
}
ReduceResult::Ignore => {
return;
}
};
let mut to_remove = SmallVec::<[usize; 2]>::new();
let ty_negated = if should_simplify_full {
ty.negate(self.db)
} else {
Type::Never // won't be used
};

if ty == element_type {
return;
for (index, element) in self.elements.iter_mut().enumerate() {
let element_type = match element.try_reduce(self.db, ty) {
ReduceResult::KeepIf(keep) => {
if !keep {
to_remove.push(index);
}
continue;
}
ReduceResult::Type(ty) => ty,
ReduceResult::CollapseToObject => {
self.collapse_to_object();
return;
}
ReduceResult::Ignore => {
return;
}
};

if Some(element_type) == bool_pair {
self.add_in_place_impl(KnownClass::Bool.to_instance(self.db), seen_aliases);
return;
}
if ty == element_type {
return;
}

if should_simplify_full && !matches!(element_type, Type::TypeAlias(_)) {
if ty.is_equivalent_to(self.db, element_type)
|| ty.is_subtype_of(self.db, element_type)
{
return;
} else if element_type.is_subtype_of(self.db, ty) {
to_remove.push(index);
} else if ty_negated.is_subtype_of(self.db, element_type) {
// We add `ty` to the union. We just checked that `~ty` is a subtype of an
// existing `element`. This also means that `~ty | ty` is a subtype of
// `element | ty`, because both elements in the first union are subtypes of
// the corresponding elements in the second union. But `~ty | ty` is just
// `object`. Since `object` is a subtype of `element | ty`, we can only
// conclude that `element | ty` must be `object` (object has no other
// supertypes). This means we can simplify the whole union to just
// `object`, since all other potential elements would also be subtypes of
// `object`.
self.collapse_to_object();
return;
}
}
}
if let Some((&first, rest)) = to_remove.split_first() {
self.elements[first] = UnionElement::Type(ty);
// We iterate in descending order to keep remaining indices valid after `swap_remove`.
for &index in rest.iter().rev() {
self.elements.swap_remove(index);
}
} else {
self.elements.push(UnionElement::Type(ty));
if Some(element_type) == bool_pair {
self.add_in_place_impl(KnownClass::Bool.to_instance(self.db), seen_aliases);
return;
}

if should_simplify_full && !matches!(element_type, Type::TypeAlias(_)) {
if ty.is_equivalent_to(self.db, element_type)
|| ty.is_subtype_of(self.db, element_type)
{
return;
} else if element_type.is_subtype_of(self.db, ty) {
to_remove.push(index);
} else if ty_negated.is_subtype_of(self.db, element_type) {
// We add `ty` to the union. We just checked that `~ty` is a subtype of an
// existing `element`. This also means that `~ty | ty` is a subtype of
// `element | ty`, because both elements in the first union are subtypes of
// the corresponding elements in the second union. But `~ty | ty` is just
// `object`. Since `object` is a subtype of `element | ty`, we can only
// conclude that `element | ty` must be `object` (object has no other
// supertypes). This means we can simplify the whole union to just
// `object`, since all other potential elements would also be subtypes of
// `object`.
self.collapse_to_object();
return;
}
}
}
if let Some((&first, rest)) = to_remove.split_first() {
self.elements[first] = UnionElement::Type(ty);
// We iterate in descending order to keep remaining indices valid after `swap_remove`.
for &index in rest.iter().rev() {
self.elements.swap_remove(index);
}
} else {
self.elements.push(UnionElement::Type(ty));
}
}

pub(crate) fn build(self) -> Type<'db> {
Expand Down
Loading