mirrors/ruff
1
0
Fork 0
mirror of https://github.com/astral-sh/ruff.git synced 2025-09-30 22:01:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

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

Browse source 
## 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 https://github.com/astral-sh/ty/issues/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.
This commit is contained in:
David Peter 2025-09-10 15:54:06 +02:00 • committed by GitHub
parent 7a75702237
commit 57d1f7132d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 83 additions and 72 deletions
Download patch file Download diff file Expand all files Collapse all files

10
crates/ty_python_semantic/resources/mdtest/union_types.md
View file

@ -118,7 +118,8 @@ def _(
```py ```py
from enum import Enum from enum import Enum
from typing import Literal from typing import Literal, Any
from ty_extensions import Intersection
class Color(Enum): class Color(Enum):
RED = "red" RED = "red"
@ -139,6 +140,13 @@ def _(
reveal_type(u4) # revealed: Literal[Color.RED, Color.GREEN] reveal_type(u4) # revealed: Literal[Color.RED, Color.GREEN]
reveal_type(u5) # revealed: Color reveal_type(u5) # revealed: Color
reveal_type(u6) # 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` ## Do not erase `Unknown`

145
crates/ty_python_semantic/src/types/builder.rs
View file

@ -444,8 +444,7 @@ impl<'db> UnionBuilder<'db> {
.filter_map(UnionElement::to_type_element) .filter_map(UnionElement::to_type_element)
.any(|ty| Type::EnumLiteral(enum_member_to_add).is_subtype_of(self.db, ty)) .any(|ty| Type::EnumLiteral(enum_member_to_add).is_subtype_of(self.db, ty))
{ {
self.elements self.push_type(Type::EnumLiteral(enum_member_to_add), seen_aliases);
.push(UnionElement::Type(Type::EnumLiteral(enum_member_to_add)));
} }
} }
// Adding `object` to a union results in `object`. // Adding `object` to a union results in `object`.
@ -453,84 +452,88 @@ impl<'db> UnionBuilder<'db> {
self.collapse_to_object(); self.collapse_to_object();
} }
_ => { _ => {
let bool_pair = if let Type::BooleanLiteral(b) = ty { self.push_type(ty, seen_aliases);
Some(Type::BooleanLiteral(!b)) }
} else { }
None }
};
// If an alias gets here, it means we aren't unpacking aliases, and we also fn push_type(&mut self, ty: Type<'db>, seen_aliases: &mut Vec<Type<'db>>) {
// shouldn't try to simplify aliases out of the union, because that will require let bool_pair = if let Type::BooleanLiteral(b) = ty {
// unpacking them. Some(Type::BooleanLiteral(!b))
let should_simplify_full = !matches!(ty, Type::TypeAlias(_)); } else {
None
};
let mut to_remove = SmallVec::<[usize; 2]>::new(); // If an alias gets here, it means we aren't unpacking aliases, and we also
let ty_negated = if should_simplify_full { // shouldn't try to simplify aliases out of the union, because that will require
ty.negate(self.db) // unpacking them.
} else { let should_simplify_full = !matches!(ty, Type::TypeAlias(_));
Type::Never // won't be used
};
for (index, element) in self.elements.iter_mut().enumerate() { let mut to_remove = SmallVec::<[usize; 2]>::new();
let element_type = match element.try_reduce(self.db, ty) { let ty_negated = if should_simplify_full {
ReduceResult::KeepIf(keep) => { ty.negate(self.db)
if !keep { } else {
to_remove.push(index); Type::Never // won't be used
} };
continue;
}
ReduceResult::Type(ty) => ty,
ReduceResult::CollapseToObject => {
self.collapse_to_object();
return;
}
ReduceResult::Ignore => {
return;
}
};
if ty == element_type { for (index, element) in self.elements.iter_mut().enumerate() {
return; let element_type = match element.try_reduce(self.db, ty) {
} ReduceResult::KeepIf(keep) => {
if !keep {
if Some(element_type) == bool_pair { to_remove.push(index);
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;
}
} }
continue;
} }
if let Some((&first, rest)) = to_remove.split_first() { ReduceResult::Type(ty) => ty,
self.elements[first] = UnionElement::Type(ty); ReduceResult::CollapseToObject => {
// We iterate in descending order to keep remaining indices valid after `swap_remove`. self.collapse_to_object();
for &index in rest.iter().rev() { return;
self.elements.swap_remove(index); }
} ReduceResult::Ignore => {
} else { return;
self.elements.push(UnionElement::Type(ty)); }
};
if ty == element_type {
return;
}
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> { pub(crate) fn build(self) -> Type<'db> {