Skip to content

Commit 57d1f71

Browse files
authored
[ty] Simplify unions of enum literals and subtypes thereof (#20324)
## Summary When adding an enum literal `E = Literal[Color.RED]` to a union which already contained a subtype of that enum literal(!), we were previously not simplifying the union correctly. My assumption is that our property tests didn't catch that earlier, because the only possible non-trivial subytpe of an enum literal that I can think of is `Any & E`. And in order for that to be detected by the property tests, it would have to randomly generate `Any & E | E` and then also compare that with `E` on the other side (in an equivalence test, or the subtyping-antisymmetry test). closes astral-sh/ty#1155 ## Test Plan * Added a regression test. * I also ran the property tests for a while, but probably not for two months worth of daily CI runs.
1 parent 7a75702 commit 57d1f71

File tree

2 files changed

+82
-71
lines changed

2 files changed

+82
-71
lines changed

crates/ty_python_semantic/resources/mdtest/union_types.md

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,8 @@ def _(
118118

119119
```py
120120
from enum import Enum
121-
from typing import Literal
121+
from typing import Literal, Any
122+
from ty_extensions import Intersection
122123

123124
class Color(Enum):
124125
RED = "red"
@@ -139,6 +140,13 @@ def _(
139140
reveal_type(u4) # revealed: Literal[Color.RED, Color.GREEN]
140141
reveal_type(u5) # revealed: Color
141142
reveal_type(u6) # revealed: Color
143+
144+
def _(
145+
u1: Intersection[Literal[Color.RED], Any] | Literal[Color.RED],
146+
u2: Literal[Color.RED] | Intersection[Literal[Color.RED], Any],
147+
):
148+
reveal_type(u1) # revealed: Literal[Color.RED]
149+
reveal_type(u2) # revealed: Literal[Color.RED]
142150
```
143151

144152
## Do not erase `Unknown`

crates/ty_python_semantic/src/types/builder.rs

Lines changed: 73 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -444,93 +444,96 @@ impl<'db> UnionBuilder<'db> {
444444
.filter_map(UnionElement::to_type_element)
445445
.any(|ty| Type::EnumLiteral(enum_member_to_add).is_subtype_of(self.db, ty))
446446
{
447-
self.elements
448-
.push(UnionElement::Type(Type::EnumLiteral(enum_member_to_add)));
447+
self.push_type(Type::EnumLiteral(enum_member_to_add), seen_aliases);
449448
}
450449
}
451450
// Adding `object` to a union results in `object`.
452451
ty if ty.is_object(self.db) => {
453452
self.collapse_to_object();
454453
}
455454
_ => {
456-
let bool_pair = if let Type::BooleanLiteral(b) = ty {
457-
Some(Type::BooleanLiteral(!b))
458-
} else {
459-
None
460-
};
455+
self.push_type(ty, seen_aliases);
456+
}
457+
}
458+
}
461459

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

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

474-
for (index, element) in self.elements.iter_mut().enumerate() {
475-
let element_type = match element.try_reduce(self.db, ty) {
476-
ReduceResult::KeepIf(keep) => {
477-
if !keep {
478-
to_remove.push(index);
479-
}
480-
continue;
481-
}
482-
ReduceResult::Type(ty) => ty,
483-
ReduceResult::CollapseToObject => {
484-
self.collapse_to_object();
485-
return;
486-
}
487-
ReduceResult::Ignore => {
488-
return;
489-
}
490-
};
472+
let mut to_remove = SmallVec::<[usize; 2]>::new();
473+
let ty_negated = if should_simplify_full {
474+
ty.negate(self.db)
475+
} else {
476+
Type::Never // won't be used
477+
};
491478

492-
if ty == element_type {
493-
return;
479+
for (index, element) in self.elements.iter_mut().enumerate() {
480+
let element_type = match element.try_reduce(self.db, ty) {
481+
ReduceResult::KeepIf(keep) => {
482+
if !keep {
483+
to_remove.push(index);
494484
}
485+
continue;
486+
}
487+
ReduceResult::Type(ty) => ty,
488+
ReduceResult::CollapseToObject => {
489+
self.collapse_to_object();
490+
return;
491+
}
492+
ReduceResult::Ignore => {
493+
return;
494+
}
495+
};
495496

496-
if Some(element_type) == bool_pair {
497-
self.add_in_place_impl(KnownClass::Bool.to_instance(self.db), seen_aliases);
498-
return;
499-
}
497+
if ty == element_type {
498+
return;
499+
}
500500

501-
if should_simplify_full && !matches!(element_type, Type::TypeAlias(_)) {
502-
if ty.is_equivalent_to(self.db, element_type)
503-
|| ty.is_subtype_of(self.db, element_type)
504-
{
505-
return;
506-
} else if element_type.is_subtype_of(self.db, ty) {
507-
to_remove.push(index);
508-
} else if ty_negated.is_subtype_of(self.db, element_type) {
509-
// We add `ty` to the union. We just checked that `~ty` is a subtype of an
510-
// existing `element`. This also means that `~ty | ty` is a subtype of
511-
// `element | ty`, because both elements in the first union are subtypes of
512-
// the corresponding elements in the second union. But `~ty | ty` is just
513-
// `object`. Since `object` is a subtype of `element | ty`, we can only
514-
// conclude that `element | ty` must be `object` (object has no other
515-
// supertypes). This means we can simplify the whole union to just
516-
// `object`, since all other potential elements would also be subtypes of
517-
// `object`.
518-
self.collapse_to_object();
519-
return;
520-
}
521-
}
522-
}
523-
if let Some((&first, rest)) = to_remove.split_first() {
524-
self.elements[first] = UnionElement::Type(ty);
525-
// We iterate in descending order to keep remaining indices valid after `swap_remove`.
526-
for &index in rest.iter().rev() {
527-
self.elements.swap_remove(index);
528-
}
529-
} else {
530-
self.elements.push(UnionElement::Type(ty));
501+
if Some(element_type) == bool_pair {
502+
self.add_in_place_impl(KnownClass::Bool.to_instance(self.db), seen_aliases);
503+
return;
504+
}
505+
506+
if should_simplify_full && !matches!(element_type, Type::TypeAlias(_)) {
507+
if ty.is_equivalent_to(self.db, element_type)
508+
|| ty.is_subtype_of(self.db, element_type)
509+
{
510+
return;
511+
} else if element_type.is_subtype_of(self.db, ty) {
512+
to_remove.push(index);
513+
} else if ty_negated.is_subtype_of(self.db, element_type) {
514+
// We add `ty` to the union. We just checked that `~ty` is a subtype of an
515+
// existing `element`. This also means that `~ty | ty` is a subtype of
516+
// `element | ty`, because both elements in the first union are subtypes of
517+
// the corresponding elements in the second union. But `~ty | ty` is just
518+
// `object`. Since `object` is a subtype of `element | ty`, we can only
519+
// conclude that `element | ty` must be `object` (object has no other
520+
// supertypes). This means we can simplify the whole union to just
521+
// `object`, since all other potential elements would also be subtypes of
522+
// `object`.
523+
self.collapse_to_object();
524+
return;
531525
}
532526
}
533527
}
528+
if let Some((&first, rest)) = to_remove.split_first() {
529+
self.elements[first] = UnionElement::Type(ty);
530+
// We iterate in descending order to keep remaining indices valid after `swap_remove`.
531+
for &index in rest.iter().rev() {
532+
self.elements.swap_remove(index);
533+
}
534+
} else {
535+
self.elements.push(UnionElement::Type(ty));
536+
}
534537
}
535538

536539
pub(crate) fn build(self) -> Type<'db> {

0 commit comments

Comments
 (0)