diff --git a/crates/pyrefly_types/src/types.rs b/crates/pyrefly_types/src/types.rs index e426b89103..15ade9ccf5 100644 --- a/crates/pyrefly_types/src/types.rs +++ b/crates/pyrefly_types/src/types.rs @@ -1680,6 +1680,8 @@ impl Type { } } } + Type::Intersect(intersect) => intersect.1.visit_toplevel_callable(f), + Type::KwCall(call) => call.return_ty.visit_toplevel_callable(f), _ => {} } } @@ -1715,6 +1717,8 @@ impl Type { } } } + Type::Intersect(intersect) => intersect.1.transform_toplevel_callable(f), + Type::KwCall(call) => call.return_ty.transform_toplevel_callable(f), _ => {} } } diff --git a/pyrefly/lib/alt/class/class_field.rs b/pyrefly/lib/alt/class/class_field.rs index a11b111150..e2e0df1991 100644 --- a/pyrefly/lib/alt/class/class_field.rs +++ b/pyrefly/lib/alt/class/class_field.rs @@ -1251,6 +1251,15 @@ fn make_bound_method_helper( } return Ok(unions(bound_methods, heap)); } + Type::Intersect(intersect) => { + let (members, fallback) = *intersect; + let bound_fallback = make_bound_method_helper(heap, obj, fallback, should_bind)?; + return Ok(heap.mk_intersect(members, bound_fallback)); + } + Type::KwCall(mut call) => { + call.return_ty = make_bound_method_helper(heap, obj, call.return_ty, should_bind)?; + return Ok(heap.mk_kw_call(*call)); + } _ => return Err(attr), }; Ok(heap.mk_bound_method(BoundMethod { obj, func })) @@ -1310,6 +1319,8 @@ pub enum DataclassMember { fn has_any_method_type(ty: &Type) -> bool { match ty { Type::Union(union) => union.members.iter().any(|t| t.has_toplevel_func_metadata()), + Type::Intersect(intersect) => intersect.1.has_toplevel_func_metadata(), + Type::KwCall(call) => has_any_method_type(&call.return_ty), _ => ty.has_toplevel_func_metadata(), } } diff --git a/pyrefly/lib/alt/function.rs b/pyrefly/lib/alt/function.rs index 6a266da643..d11ad9b7bd 100644 --- a/pyrefly/lib/alt/function.rs +++ b/pyrefly/lib/alt/function.rs @@ -1671,6 +1671,25 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { { original_decoratee.clone() } + Type::ClassType(cls) if cls.has_qname("functools", "_lru_cache_wrapper") => { + // Keep the original callable as a real intersection member so simplification + // cannot collapse the type to just the wrapper. + self.heap.mk_intersect( + vec![self.heap.mk_class_type(cls), original_decoratee.clone()], + original_decoratee.clone(), + ) + } + Type::KwCall(mut call) if matches!(&call.return_ty, Type::ClassType(cls) if cls.has_qname("functools", "_lru_cache_wrapper")) => { + if original_decoratee.property_metadata().is_some() { + original_decoratee.clone() + } else { + call.return_ty = self.heap.mk_intersect( + vec![call.return_ty.clone(), original_decoratee.clone()], + original_decoratee.clone(), + ); + self.heap.mk_kw_call(*call) + } + } // Heuristic for union-typed decorators (e.g. dual-use decorators that // can be applied with or without parentheses like @d or @d(flag)): // @@ -1916,6 +1935,15 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { errors: &ErrorCollector, ) -> Vec1<(TextRange, OverloadType)> { ts.mapped(|(range, t, metadata)| { + let display_t = t.clone(); + let mut t = t; + let t = loop { + match t { + Type::Intersect(intersect) => t = intersect.1, + Type::KwCall(call) => t = call.return_ty, + t => break t, + } + }; ( range, match t { @@ -1967,7 +1995,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { format!( "`{}` has type `{}` after decorator application, which is not callable", func, - self.for_display(t) + self.for_display(display_t) ), ); OverloadType::Function(Function { diff --git a/pyrefly/lib/test/functools/cache.rs b/pyrefly/lib/test/functools/cache.rs new file mode 100644 index 0000000000..40cb8ea2ed --- /dev/null +++ b/pyrefly/lib/test/functools/cache.rs @@ -0,0 +1,125 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +//! `functools.cache` and `functools.lru_cache` preserve the wrapped function's call signature. + +use crate::functools_testcase; +use crate::test::util::TestEnv; +use crate::testcase; + +fn cached_overload_env() -> TestEnv { + let mut env = TestEnv::new(); + env.add_with_path( + "pkg", + "pkg/__init__.pyi", + r#" +from functools import lru_cache +from typing import Literal, Self, overload + +class HasProps: + @overload + @classmethod + @lru_cache + def properties(cls: type[Self], *, _with_props: Literal[False] = False) -> set[str]: ... + @overload + @classmethod + @lru_cache + def properties(cls: type[Self], *, _with_props: Literal[True] = True) -> dict[str, object]: ... +"#, + ); + env +} + +functools_testcase!( + test_cache_preserves_function_signature, + r#" +from functools import cache +from typing import assert_type + +@cache +def f() -> int: + return 0 + +assert_type(f(), int) +f(1, 2, 3) # E: Expected 0 positional arguments, got 3 +f.cache_clear() +assert_type(f.cache_info().hits, int) +"#, +); + +functools_testcase!( + test_lru_cache_preserves_function_signature, + r#" +from functools import lru_cache +from typing import assert_type + +@lru_cache +def f(x: int, y: str = "") -> bool: + return True + +assert_type(f(1), bool) +assert_type(f(1, "x"), bool) +f("x") # E: Argument `Literal['x']` is not assignable to parameter `x` with type `int` +f(1, y=2) # E: Argument `Literal[2]` is not assignable to parameter `y` with type `str` +f.cache_info() +"#, +); + +functools_testcase!( + test_lru_cache_call_preserves_function_signature, + r#" +from functools import lru_cache +from typing import assert_type + +@lru_cache(maxsize=None) +def f(x: int) -> str: + return "" + +assert_type(f(1), str) +f("x") # E: Argument `Literal['x']` is not assignable to parameter `x` with type `int` +f.cache_clear() +"#, +); + +functools_testcase!( + test_lru_cache_method_binds_self, + r#" +from functools import cache, lru_cache +from typing import assert_type + +class C: + @cache + def cache_method(self, x: int) -> str: + return "" + + @lru_cache(maxsize=4) + def lru_method(self, x: int) -> str: + return "" + +c = C() +assert_type(C.cache_method(c, 1), str) +assert_type(c.cache_method(1), str) +c.cache_method("x") # E: Argument `Literal['x']` is not assignable to parameter `x` with type `int` +c.cache_method.cache_info() +assert_type(C.lru_method(c, 1), str) +assert_type(c.lru_method(1), str) +c.lru_method("x") # E: Argument `Literal['x']` is not assignable to parameter `x` with type `int` +c.lru_method.cache_clear() +"#, +); + +testcase!( + test_lru_cache_preserves_overload_signatures_in_stubs, + cached_overload_env(), + r#" +from pkg import HasProps +from typing import assert_type + +assert_type(HasProps.properties(), set[str]) +assert_type(HasProps.properties(_with_props=True), dict[str, object]) +"#, +); diff --git a/pyrefly/lib/test/functools/mod.rs b/pyrefly/lib/test/functools/mod.rs index 12436e678e..48f4aeab29 100644 --- a/pyrefly/lib/test/functools/mod.rs +++ b/pyrefly/lib/test/functools/mod.rs @@ -6,6 +6,7 @@ */ #![cfg(test)] +mod cache; mod partial; mod partial_edge; mod partial_generics; diff --git a/pyrefly/lib/test/pysa/call_graph.rs b/pyrefly/lib/test/pysa/call_graph.rs index f72520add4..1df55db87e 100644 --- a/pyrefly/lib/test/pysa/call_graph.rs +++ b/pyrefly/lib/test/pysa/call_graph.rs @@ -6282,7 +6282,10 @@ def foo(): vec![ ( "4:3-6:13", - define_callees(vec![create_call_target("test.inner", TargetType::Function)]), + define_callees(vec![ + create_call_target("test.inner", TargetType::Function) + .with_return_type(ScalarTypeProperties::int()), + ]), ), ( "7:3-7:10",