[ty] Handle decorators which return unions of Callables (#20858)

## Summary

If a function is decorated with a decorator that returns a union of
`Callable`s, also treat it as a union of function-like `Callable`s.

Labeling as `internal`, since the previous change has not been released
yet.

## Test Plan

New regression test.
This commit is contained in:
David Peter 2025-10-14 11:47:50 +02:00 committed by GitHub
parent c69fa75cd5
commit ac2c530377
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 43 additions and 9 deletions

View file

@ -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

View file

@ -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
}