mirrors/ruff
1
0
Fork 0
mirror of https://github.com/astral-sh/ruff.git synced 2025-09-30 13:51:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

[ty] Introduce UnionType::try_from_elements and UnionType::try_map (#18911)

Browse source
This commit is contained in:
Alex Waygood 2025-06-24 13:09:02 +01:00 • committed by GitHub
parent 27eee5a1a8
commit 237a5821ba
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 67 additions and 57 deletions
Download patch file Download diff file Expand all files Collapse all files

58
crates/ty_python_semantic/src/types.rs
View file

@ -3484,11 +3484,7 @@ impl<'db> Type<'db> {
Type::IntLiteral(value) => (value >= 0).then_some(ty),
Type::BooleanLiteral(value) => Some(Type::IntLiteral(value.into())),
Type::Union(union) => {
let mut builder = UnionBuilder::new(db);
for element in union.elements(db) {
builder = builder.add(non_negative_int_literal(db, *element)?);
}
Some(builder.build())
union.try_map(db, |element| non_negative_int_literal(db, *element))
}
_ => None,
}
@ -4801,13 +4797,7 @@ impl<'db> Type<'db> {
Type::ClassLiteral(class) => Some(Type::instance(db, class.default_specialization(db))),
Type::GenericAlias(alias) => Some(Type::instance(db, ClassType::from(*alias))),
Type::SubclassOf(subclass_of_ty) => Some(subclass_of_ty.to_instance(db)),
Type::Union(union) => {
let mut builder = UnionBuilder::new(db);
for element in union.elements(db) {
builder = builder.add(element.to_instance(db)?);
}
Some(builder.build())
}
Type::Union(union) => union.to_instance(db),
// If there is no bound or constraints on a typevar `T`, `T: object` implicitly, which
// has no instance type. Otherwise, synthesize a typevar with bound or constraints
// mapped through `to_instance`.
@ -4817,11 +4807,9 @@ impl<'db> Type<'db> {
TypeVarBoundOrConstraints::UpperBound(upper_bound.to_instance(db)?)
}
TypeVarBoundOrConstraints::Constraints(constraints) => {
let mut builder = UnionBuilder::new(db);
for constraint in constraints.elements(db) {
builder = builder.add(constraint.to_instance(db)?);
}
TypeVarBoundOrConstraints::Constraints(builder.build().into_union()?)
TypeVarBoundOrConstraints::Constraints(
constraints.to_instance(db)?.into_union()?,
)
}
};
Some(Type::TypeVar(TypeVarInstance::new(
@ -7640,6 +7628,23 @@ impl<'db> UnionType<'db> {
.build()
}
/// A fallible version of [`UnionType::from_elements`].
///
/// If all items in `elements` are `Some()`, the result of unioning all elements is returned.
/// As soon as a `None` element in the iterable is encountered,
/// the function short-circuits and returns `None`.
pub(crate) fn try_from_elements<I, T>(db: &'db dyn Db, elements: I) -> Option<Type<'db>>
where
I: IntoIterator<Item = Option<T>>,
T: Into<Type<'db>>,
{
let mut builder = UnionBuilder::new(db);
for element in elements {
builder = builder.add(element?.into());
}
Some(builder.build())
}
/// Apply a transformation function to all elements of the union,
/// and create a new union from the resulting set of types.
pub fn map(
@ -7650,6 +7655,25 @@ impl<'db> UnionType<'db> {
Self::from_elements(db, self.elements(db).iter().map(transform_fn))
}
/// A fallible version of [`UnionType::map`].
///
/// For each element in `self`, `transform_fn` is called on that element.
/// If `transform_fn` returns `Some()` for all elements in `self`,
/// the result of unioning all transformed elements is returned.
/// As soon as `transform_fn` returns `None` for an element, however,
/// the function short-circuits and returns `None`.
pub(crate) fn try_map(
self,
db: &'db dyn Db,
transform_fn: impl FnMut(&Type<'db>) -> Option<Type<'db>>,
) -> Option<Type<'db>> {
Self::try_from_elements(db, self.elements(db).iter().map(transform_fn))
}
pub(crate) fn to_instance(self, db: &'db dyn Db) -> Option<Type<'db>> {
self.try_map(db, |element| element.to_instance(db))
}
pub(crate) fn filter(
self,
db: &'db dyn Db,

46
crates/ty_python_semantic/src/types/infer.rs
View file

@ -6770,34 +6770,24 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
}
match (left_ty, right_ty, op) {
(Type::Union(lhs_union), rhs, _) => {
let mut union = UnionBuilder::new(self.db());
for lhs in lhs_union.elements(self.db()) {
let result = self.infer_binary_expression_type(
node,
emitted_division_by_zero_diagnostic,
*lhs,
rhs,
op,
)?;
union = union.add(result);
}
Some(union.build())
}
(lhs, Type::Union(rhs_union), _) => {
let mut union = UnionBuilder::new(self.db());
for rhs in rhs_union.elements(self.db()) {
let result = self.infer_binary_expression_type(
node,
emitted_division_by_zero_diagnostic,
lhs,
*rhs,
op,
)?;
union = union.add(result);
}
Some(union.build())
}
(Type::Union(lhs_union), rhs, _) => lhs_union.try_map(self.db(), |lhs_element| {
self.infer_binary_expression_type(
node,
emitted_division_by_zero_diagnostic,
*lhs_element,
rhs,
op,
)
}),
(lhs, Type::Union(rhs_union), _) => rhs_union.try_map(self.db(), |rhs_element| {
self.infer_binary_expression_type(
node,
emitted_division_by_zero_diagnostic,
lhs,
*rhs_element,
op,
)
}),
// Non-todo Anys take precedence over Todos (as if we fix this `Todo` in the future,
// the result would then become Any or Unknown, respectively).

20
crates/ty_python_semantic/src/types/narrow.rs
View file

@ -177,13 +177,13 @@ impl ClassInfoConstraintFunction {
};
match classinfo {
Type::Tuple(tuple) => {
let mut builder = UnionBuilder::new(db);
for element in tuple.tuple(db).all_elements() {
builder = builder.add(self.generate_constraint(db, element)?);
}
Some(builder.build())
}
Type::Tuple(tuple) => UnionType::try_from_elements(
db,
tuple
.tuple(db)
.all_elements()
.map(|element| self.generate_constraint(db, element)),
),
Type::ClassLiteral(class_literal) => {
// At runtime (on Python 3.11+), this will return `True` for classes that actually
// do inherit `typing.Any` and `False` otherwise. We could accurately model that?
@ -214,11 +214,7 @@ impl ClassInfoConstraintFunction {
}
}
Type::Union(union) => {
let mut builder = UnionBuilder::new(db);
for element in union.elements(db) {
builder = builder.add(self.generate_constraint(db, *element)?);
}
Some(builder.build())
union.try_map(db, |element| self.generate_constraint(db, *element))
}
Type::TypeVar(type_var) => match type_var.bound_or_constraints(db)? {
TypeVarBoundOrConstraints::UpperBound(bound) => self.generate_constraint(db, bound),