[ty] Handle decorators which return unions of Callables

crates/ty_python_semantic/resources/mdtest/call/callables_as_descriptors.md

@@ -145,22 +145,38 @@ class C2:
 C2().method_decorated(1)
 ```
 
+And with unions of `Callable` types:
+
+```py
+from typing import Callable
+
+def expand(f: Callable[[C3, int], int]) -> Callable[[C3, int], int] | Callable[[C3, int], str]:
+    raise NotImplementedError
+
+class C3:
+    @expand
+    def method_decorated(self, x: int) -> int:
+        return x
+
+reveal_type(C3().method_decorated(1))  # revealed: int | str
+```
+
 Note that we currently only apply this heuristic when calling a function such as `memoize` via the
 decorator syntax. This is inconsistent, because the above *should* be equivalent to the following,
 but here we emit errors:
 
 ```py
-def memoize3(f: Callable[[C3, int], str]) -> Callable[[C3, int], str]:
+def memoize3(f: Callable[[C4, int], str]) -> Callable[[C4, int], str]:
     raise NotImplementedError
 
-class C3:
+class C4:
     def method(self, x: int) -> str:
         return str(x)
     method_decorated = memoize3(method)
 
 # error: [missing-argument]
 # error: [invalid-argument-type]
-C3().method_decorated(1)
+C4().method_decorated(1)
 ```
 
 The reason for this is that the heuristic is problematic. We don't *know* that the `Callable` in the

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -2199,19 +2199,37 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 .map(|bindings| bindings.return_type(self.db()))
             {
                 Ok(return_ty) => {
+                    fn into_function_like_callable<'d>(
+                        db: &'d dyn Db,
+                        ty: Type<'d>,
+                    ) -> Option<Type<'d>> {
+                        match ty {
+                            Type::Callable(callable) => Some(Type::Callable(CallableType::new(
+                                db,
+                                callable.signatures(db),
+                                true,
+                            ))),
+                            Type::Union(union) => union
+                                .try_map(db, |element| into_function_like_callable(db, *element)),
+                            // Intersections are currently not handled here because that would require
+                            // the decorator to be explicitly annotated as returning an intersection.
+                            _ => None,
+                        }
+                    }
+
                     let is_input_function_like = inferred_ty
                         .try_upcast_to_callable(self.db())
                         .and_then(Type::as_callable)
                         .is_some_and(|callable| callable.is_function_like(self.db()));
-                    if is_input_function_like && let Some(callable_type) = return_ty.as_callable() {
+
+                    if is_input_function_like
+                        && let Some(return_ty_function_like) =
+                            into_function_like_callable(self.db(), return_ty)
+                    {
                         // When a method on a class is decorated with a function that returns a `Callable`, assume that
                         // the returned callable is also function-like. See "Decorating a method with a `Callable`-typed
                         // decorator" in `callables_as_descriptors.md` for the extended explanation.
-                        Type::Callable(CallableType::new(
-                            self.db(),
-                            callable_type.signatures(self.db()),
-                            true,
-                        ))
+                        return_ty_function_like
                     } else {
                         return_ty
                     }