From 57d1f7132dbac82e07f22560dbedba7cff1dadc4 Mon Sep 17 00:00:00 2001 From: David Peter Date: Wed, 10 Sep 2025 15:54:06 +0200 Subject: [PATCH] [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 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. --- .../resources/mdtest/union_types.md | 10 +- .../ty_python_semantic/src/types/builder.rs | 145 +++++++++--------- 2 files changed, 83 insertions(+), 72 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/union_types.md b/crates/ty_python_semantic/resources/mdtest/union_types.md index 52a77028c7..3974b98993 100644 --- a/crates/ty_python_semantic/resources/mdtest/union_types.md +++ b/crates/ty_python_semantic/resources/mdtest/union_types.md @@ -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" @@ -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` diff --git a/crates/ty_python_semantic/src/types/builder.rs b/crates/ty_python_semantic/src/types/builder.rs index 97bcf20f9b..16b28e5fe8 100644 --- a/crates/ty_python_semantic/src/types/builder.rs +++ b/crates/ty_python_semantic/src/types/builder.rs @@ -444,8 +444,7 @@ 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))); + self.push_type(Type::EnumLiteral(enum_member_to_add), seen_aliases); } } // Adding `object` to a union results in `object`. @@ -453,84 +452,88 @@ impl<'db> UnionBuilder<'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>) { + 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; - } - - 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; - } + 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; } - 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)); + ReduceResult::Type(ty) => ty, + ReduceResult::CollapseToObject => { + self.collapse_to_object(); + return; + } + ReduceResult::Ignore => { + return; + } + }; + + 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> {