diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index a1c9076b9b..eafe6d4f67 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -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(db: &'db dyn Db, elements: I) -> Option> + where + I: IntoIterator>, + T: Into>, + { + 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>, + ) -> Option> { + Self::try_from_elements(db, self.elements(db).iter().map(transform_fn)) + } + + pub(crate) fn to_instance(self, db: &'db dyn Db) -> Option> { + self.try_map(db, |element| element.to_instance(db)) + } + pub(crate) fn filter( self, db: &'db dyn Db, diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index 997e3a6eed..c5d08f5a7e 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -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). diff --git a/crates/ty_python_semantic/src/types/narrow.rs b/crates/ty_python_semantic/src/types/narrow.rs index 91296446f7..ca03a38f41 100644 --- a/crates/ty_python_semantic/src/types/narrow.rs +++ b/crates/ty_python_semantic/src/types/narrow.rs @@ -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),