Override more pure-wrapper decorators as safe

Summary: More stubs + more tests

Reviewed By: martindemello

Differential Revision: D107561779

fbshipit-source-id: e3875869b34a3a364e057e66b1d002953e4185e1

Author: brittanyrey

resources/stubs/stdlib/functools.pyi

@@ -148,7 +148,7 @@ else:
         updated: Iterable[str] = ("__dict__",),
     ) -> _Wrapper[_PWrapped, _RWrapped]: ...
 
-def total_ordering(cls: type[_T]) -> type[_T]: ...
+def total_ordering(cls: type[_T]) -> type[_T]: no_effects()
 def cmp_to_key(mycmp: Callable[[_T, _T], int]) -> Callable[[_T], SupportsAllComparisons]: ...
 @disjoint_base
 class partial(Generic[_T]):

src/manual_override.rs

@@ -26,8 +26,14 @@ const SAFE_FUNCTIONS_ARRAY: &[&str] = &[
     // Decorators for caching values.  These are akin to functools.cache().
     "f3.utils.decorators.cache",
     "libfb.py.decorators.lazy_property",
+    "libfb.py.decorators.memoize_fast",
+    "libfb.py.decorators.memoize_fast_0",
     "libfb.py.decorators.memoize_forever",
     "libfb.py.decorators.memoize_timed",
+    "libfb.py.memoize.memoize_fast",
+    "libfb.py.memoize.memoize_fast_0",
+    "psutil._common.memoize",
+    "psutil._common.memoize_when_activated",
     // Functions that use the ABCMeta registry, which is hard to analyze but it is safe.
     "collections.abc.Mapping.register",
     "collections.abc.Sequence.register",
@@ -106,10 +112,14 @@ const SAFE_FUNCTIONS_ARRAY: &[&str] = &[
     "dns.name.IDNA2003Codec",
     "dns.name.IDNA2008Codec",
     "dns.name.Name",
+    "google.auth._helpers.copy_docstring",
     // Pure-wrapper deprecation decorators. At decoration time they only build a
     // closure that emits a DeprecationWarning at call time; no module-level
     // registry, no I/O. Same shape as langchain_core._api.deprecation.deprecated.
     "markdown.util.deprecated",
+    // Pure functools.wraps wrapper. All tracing work happens inside the async
+    // wrapper at call time, never at decoration time.
+    "model_context_protocol.common.decorators.mcp_client_connect_decorator.mcp_client_connect",
     "nltk.internals.deprecated",
     "sympy.utilities.decorator.deprecated",
     // sandcastle test gate: at decoration time it either returns the function

tests/stubs.rs

@@ -182,6 +182,99 @@ def f(new=1):
         check(code);
     }
 
+    #[test]
+    fn test_libfb_memoize_fast_safe() {
+        // Regression test for the libfb memoize_fast / memoize_fast_0 overrides
+        // in manual_override.rs. Both build a cache-holding closure and return it
+        // at decoration time; the wrapped function only runs on first call.
+        let code = r#"
+from libfb.py.decorators import memoize_fast, memoize_fast_0
+
+@memoize_fast
+def f(key):
+    return key
+
+@memoize_fast_0
+def g():
+    return 1
+"#;
+        check(code);
+    }
+
+    #[test]
+    fn test_psutil_memoize_safe() {
+        // Regression test for the psutil memoize / memoize_when_activated overrides
+        // in manual_override.rs. Both wrap via functools.wraps and cache at call
+        // time; nothing mutates module state at decoration time.
+        let code = r#"
+from psutil._common import memoize, memoize_when_activated
+
+@memoize
+def f(x):
+    return x
+
+@memoize_when_activated
+def g(self):
+    return 1
+"#;
+        check(code);
+    }
+
+    #[test]
+    fn test_google_auth_copy_docstring_safe() {
+        // Regression test for the google.auth._helpers.copy_docstring override in
+        // manual_override.rs. It returns a closure that copies a docstring onto
+        // the decorated method only, with no module-level side effects.
+        let code = r#"
+from google.auth._helpers import copy_docstring
+
+class Base:
+    def m(self):
+        "doc"
+
+class C:
+    @copy_docstring(Base)
+    def m(self):
+        return 1
+"#;
+        check(code);
+    }
+
+    #[test]
+    fn test_functools_total_ordering_safe() {
+        // Regression test for the no_effects() annotation on functools.total_ordering
+        // in the bundled stdlib stub. It only injects comparison methods onto the
+        // decorated class, with no module-scope side effects.
+        let code = r#"
+from functools import total_ordering
+
+@total_ordering
+class C:
+    def __eq__(self, other):
+        return True
+
+    def __lt__(self, other):
+        return False
+"#;
+        check(code);
+    }
+
+    #[test]
+    fn test_mcp_client_connect_safe() {
+        // Regression test for the mcp_client_connect override in manual_override.rs.
+        // It is a pure functools.wraps wrapper; all tracing work happens inside the
+        // async wrapper at call time, never at decoration time.
+        let code = r#"
+from model_context_protocol.common.decorators.mcp_client_connect_decorator import mcp_client_connect
+
+class C:
+    @mcp_client_connect
+    async def connect_to_server(self, config):
+        return None
+"#;
+        check(code);
+    }
+
     #[test]
     fn test_pydantic_stub_class_body_calls_safe() {
         // Regression test for the bundled pydantic/__init__.pyi stub.