diff --git a/crates/ty_python_semantic/resources/mdtest/call/callables_as_descriptors.md b/crates/ty_python_semantic/resources/mdtest/call/callables_as_descriptors.md index cee7df8eca..2adfd228c5 100644 --- a/crates/ty_python_semantic/resources/mdtest/call/callables_as_descriptors.md +++ b/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 diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index b09b3f0196..bae41b14b4 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/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> { + 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 }