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

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);
+                    union.try_map(db, |element| non_negative_int_literal(db, *element))
-                    for element in union.elements(db) {
-                        builder = builder.add(non_negative_int_literal(db, *element)?);
-                    }
-                    Some(builder.build())
                 }
                 _ => 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) => {
+            Type::Union(union) => union.to_instance(db),
-                let mut builder = UnionBuilder::new(db);
-                for element in union.elements(db) {
-                    builder = builder.add(element.to_instance(db)?);
-                }
-                Some(builder.build())
-            }
             // 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);
+                        TypeVarBoundOrConstraints::Constraints(
-                        for constraint in constraints.elements(db) {
+                            constraints.to_instance(db)?.into_union()?,
-                            builder = builder.add(constraint.to_instance(db)?);
+                        )
-                        }
-                        TypeVarBoundOrConstraints::Constraints(builder.build().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,

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, _) => {
+            (Type::Union(lhs_union), rhs, _) => lhs_union.try_map(self.db(), |lhs_element| {
-                let mut union = UnionBuilder::new(self.db());
+                self.infer_binary_expression_type(
-                for lhs in lhs_union.elements(self.db()) {
+                    node,
-                    let result = self.infer_binary_expression_type(
+                    emitted_division_by_zero_diagnostic,
-                        node,
+                    *lhs_element,
-                        emitted_division_by_zero_diagnostic,
+                    rhs,
-                        *lhs,
+                    op,
-                        rhs,
+                )
-                        op,
+            }),
-                    )?;
+            (lhs, Type::Union(rhs_union), _) => rhs_union.try_map(self.db(), |rhs_element| {
-                    union = union.add(result);
+                self.infer_binary_expression_type(
-                }
+                    node,
-                Some(union.build())
+                    emitted_division_by_zero_diagnostic,
-            }
+                    lhs,
-            (lhs, Type::Union(rhs_union), _) => {
+                    *rhs_element,
-                let mut union = UnionBuilder::new(self.db());
+                    op,
-                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())
-            }
 
             // 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).

crates/ty_python_semantic/src/types/narrow.rs

@@ -177,13 +177,13 @@ impl ClassInfoConstraintFunction {
         };
 
         match classinfo {
-            Type::Tuple(tuple) => {
+            Type::Tuple(tuple) => UnionType::try_from_elements(
-                let mut builder = UnionBuilder::new(db);
+                db,
-                for element in tuple.tuple(db).all_elements() {
+                tuple
-                    builder = builder.add(self.generate_constraint(db, element)?);
+                    .tuple(db)
-                }
+                    .all_elements()
-                Some(builder.build())
+                    .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);
+                union.try_map(db, |element| self.generate_constraint(db, *element))
-                for element in union.elements(db) {
-                    builder = builder.add(self.generate_constraint(db, *element)?);
-                }
-                Some(builder.build())
             }
             Type::TypeVar(type_var) => match type_var.bound_or_constraints(db)? {
                 TypeVarBoundOrConstraints::UpperBound(bound) => self.generate_constraint(db, bound),