Use distribute_over_union helper for call exprs

Summary: Another small cleanup. Reviewed By: ndmitchell Differential Revision: D66490435 fbshipit-source-id: 6970fc7787729f3f855cd521d9254a74e85351b0
facebook · Nov 26, 2024 · 6d55815 · 6d55815
1 parent e5054ef
commit 6d55815
Showing 1 changed file with 19 additions and 17 deletions.
diff --git a/pyre2/pyre2/bin/alt/expr.rs b/pyre2/pyre2/bin/alt/expr.rs
@@ -842,20 +842,6 @@ impl<'a> AnswersSolver<'a> {
             }
             Expr::Call(x) => {
                 let ty_fun = self.expr_infer(&x.func);
-                let func_range = x.func.range();
-                let check_call = |ty: Type| -> Type {
-                    self.as_callable_or_error(ty.clone(), CallStyle::FreeForm, func_range, |c| {
-                        self.call_infer(
-                            c,
-                            &x.arguments.args,
-                            &x.arguments.keywords,
-                            &|arg, hint| {
-                                self.expr(arg, hint);
-                            },
-                            func_range,
-                        )
-                    })
-                };
                 if TypeVar::is_ctor(&ty_fun) {
                     Type::type_form(self.tyvar_from_arguments(&x.arguments).to_type())
                 } else if TypeVarTuple::is_ctor(&ty_fun)
@@ -866,10 +852,26 @@ impl<'a> AnswersSolver<'a> {
                     && let Some(name) = arguments_one_string(&x.arguments)
                 {
                     Type::type_form(ParamSpec::new(name, self.module_info().dupe()).to_type())
-                } else if let Type::Union(members) = ty_fun {
-                    self.unions(&members.into_iter().map(check_call).collect::<Vec<_>>())
                 } else {
-                    check_call(ty_fun.clone())
+                    let func_range = x.func.range();
+                    self.distribute_over_union(&ty_fun, |ty| {
+                        self.as_callable_or_error(
+                            ty.clone(),
+                            CallStyle::FreeForm,
+                            func_range,
+                            |c| {
+                                self.call_infer(
+                                    c,
+                                    &x.arguments.args,
+                                    &x.arguments.keywords,
+                                    &|arg, hint| {
+                                        self.expr(arg, hint);
+                                    },
+                                    func_range,
+                                )
+                            },
+                        )
+                    })
                 }
             }
             Expr::FString(x) => {