Skip to content

Commit ac2c530

Browse files
authored
[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.
1 parent c69fa75 commit ac2c530

File tree

2 files changed

+43
-9
lines changed

2 files changed

+43
-9
lines changed

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

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -145,22 +145,38 @@ class C2:
145145
C2().method_decorated(1)
146146
```
147147

148+
And with unions of `Callable` types:
149+
150+
```py
151+
from typing import Callable
152+
153+
def expand(f: Callable[[C3, int], int]) -> Callable[[C3, int], int] | Callable[[C3, int], str]:
154+
raise NotImplementedError
155+
156+
class C3:
157+
@expand
158+
def method_decorated(self, x: int) -> int:
159+
return x
160+
161+
reveal_type(C3().method_decorated(1)) # revealed: int | str
162+
```
163+
148164
Note that we currently only apply this heuristic when calling a function such as `memoize` via the
149165
decorator syntax. This is inconsistent, because the above *should* be equivalent to the following,
150166
but here we emit errors:
151167

152168
```py
153-
def memoize3(f: Callable[[C3, int], str]) -> Callable[[C3, int], str]:
169+
def memoize3(f: Callable[[C4, int], str]) -> Callable[[C4, int], str]:
154170
raise NotImplementedError
155171

156-
class C3:
172+
class C4:
157173
def method(self, x: int) -> str:
158174
return str(x)
159175
method_decorated = memoize3(method)
160176

161177
# error: [missing-argument]
162178
# error: [invalid-argument-type]
163-
C3().method_decorated(1)
179+
C4().method_decorated(1)
164180
```
165181

166182
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

Lines changed: 24 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2199,19 +2199,37 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
21992199
.map(|bindings| bindings.return_type(self.db()))
22002200
{
22012201
Ok(return_ty) => {
2202+
fn into_function_like_callable<'d>(
2203+
db: &'d dyn Db,
2204+
ty: Type<'d>,
2205+
) -> Option<Type<'d>> {
2206+
match ty {
2207+
Type::Callable(callable) => Some(Type::Callable(CallableType::new(
2208+
db,
2209+
callable.signatures(db),
2210+
true,
2211+
))),
2212+
Type::Union(union) => union
2213+
.try_map(db, |element| into_function_like_callable(db, *element)),
2214+
// Intersections are currently not handled here because that would require
2215+
// the decorator to be explicitly annotated as returning an intersection.
2216+
_ => None,
2217+
}
2218+
}
2219+
22022220
let is_input_function_like = inferred_ty
22032221
.try_upcast_to_callable(self.db())
22042222
.and_then(Type::as_callable)
22052223
.is_some_and(|callable| callable.is_function_like(self.db()));
2206-
if is_input_function_like && let Some(callable_type) = return_ty.as_callable() {
2224+
2225+
if is_input_function_like
2226+
&& let Some(return_ty_function_like) =
2227+
into_function_like_callable(self.db(), return_ty)
2228+
{
22072229
// When a method on a class is decorated with a function that returns a `Callable`, assume that
22082230
// the returned callable is also function-like. See "Decorating a method with a `Callable`-typed
22092231
// decorator" in `callables_as_descriptors.md` for the extended explanation.
2210-
Type::Callable(CallableType::new(
2211-
self.db(),
2212-
callable_type.signatures(self.db()),
2213-
true,
2214-
))
2232+
return_ty_function_like
22152233
} else {
22162234
return_ty
22172235
}

0 commit comments

Comments
 (0)