fix Typed functools.Cache #4010

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),
             _ => {}
         }
     }

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(),
     }
 }

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 {

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])
+"#,
+);

pyrefly/lib/test/functools/mod.rs

@@ -6,6 +6,7 @@
  */
 
 #![cfg(test)]
+mod cache;
 mod partial;
 mod partial_edge;
 mod partial_generics;

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",