From 2b57c610ddf75ec0e87e6edabd455e998a0371de Mon Sep 17 00:00:00 2001
From: Matthew Roeschke <10647082+mroeschke@users.noreply.github.com>
Date: Thu, 22 Feb 2024 13:50:04 -1000
Subject: [PATCH] Ensure slow private attrs are maybe proxies (#14380)

Expected pandas test failures:

>  tests/indexing/test_indexing.py

Due to this PR, it appears an `assert something._values is something_else` fails more after this PR since `._values` wraps objects in an proxy object now (a known failure mode)

> tests/series/indexing/test_setitem.py

Runs into the issue where a test set up calls `proxy._values[key] = something` using a pandas helper function that isn't proxying correctly

Authors:
  - Matthew Roeschke (https://github.com/mroeschke)
  - GALI PREM SAGAR (https://github.com/galipremsagar)

Approvers:
  - Bradley Dice (https://github.com/bdice)
  - GALI PREM SAGAR (https://github.com/galipremsagar)

URL: https://github.com/rapidsai/cudf/pull/14380
---
 python/cudf/cudf/pandas/fast_slow_proxy.py    | 19 +++++++++++-
 .../cudf_pandas_tests/test_cudf_pandas.py     |  7 +++++
 .../cudf_pandas_tests/test_fast_slow_proxy.py | 31 ++++++++++++++++++-
 3 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/python/cudf/cudf/pandas/fast_slow_proxy.py b/python/cudf/cudf/pandas/fast_slow_proxy.py
index d132116af61..a2b14e0c3aa 100644
--- a/python/cudf/cudf/pandas/fast_slow_proxy.py
+++ b/python/cudf/cudf/pandas/fast_slow_proxy.py
@@ -572,7 +572,24 @@ def __getattr__(self, name: str) -> Any:
             _raise_attribute_error(self.__class__.__name__, name)
         if name.startswith("_"):
             # private attributes always come from `._fsproxy_slow`:
-            return getattr(self._fsproxy_slow, name)
+            obj = getattr(self._fsproxy_slow, name)
+            if name.startswith("__array"):
+                # TODO: numpy methods raise when given proxy ndarray objects
+                # https://numpy.org/doc/stable/reference/arrays.classes.html#special-attributes-and-methods  # noqa:E501
+                return obj
+
+            if not _is_function_or_method(obj):
+                return _maybe_wrap_result(
+                    obj, getattr, self._fsproxy_slow, name
+                )
+
+            @functools.wraps(obj)
+            def _wrapped_private_slow(*args, **kwargs):
+                slow_args, slow_kwargs = _slow_arg(args), _slow_arg(kwargs)
+                result = obj(*slow_args, **slow_kwargs)
+                return _maybe_wrap_result(result, obj, *args, **kwargs)
+
+            return _wrapped_private_slow
         attr = _FastSlowAttribute(name)
         return attr.__get__(self)
 
diff --git a/python/cudf/cudf_pandas_tests/test_cudf_pandas.py b/python/cudf/cudf_pandas_tests/test_cudf_pandas.py
index ab4742549f8..0386ec434da 100644
--- a/python/cudf/cudf_pandas_tests/test_cudf_pandas.py
+++ b/python/cudf/cudf_pandas_tests/test_cudf_pandas.py
@@ -1078,6 +1078,13 @@ def test_dataframe_query():
     tm.assert_equal(actual, expected)
 
 
+def test_private_method_result_wrapped():
+    xoffset = xpd.offsets.Day()
+    dt = datetime.datetime(2020, 1, 1)
+    result = xoffset._apply(dt)
+    assert isinstance(result, xpd.Timestamp)
+
+
 def test_numpy_var():
     np.random.seed(42)
     data = np.random.rand(1000)
diff --git a/python/cudf/cudf_pandas_tests/test_fast_slow_proxy.py b/python/cudf/cudf_pandas_tests/test_fast_slow_proxy.py
index b964dfde4ed..631ad2f37b2 100644
--- a/python/cudf/cudf_pandas_tests/test_fast_slow_proxy.py
+++ b/python/cudf/cudf_pandas_tests/test_fast_slow_proxy.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2023-2024 NVIDIA CORPORATION & AFFILIATES.
 # All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
@@ -445,6 +445,35 @@ def __radd__(self, other):
     assert BarProxy() + Foo() == "sum"
 
 
+def test_slow_attr_still_proxy():
+    class A:
+        pass
+
+    class B:
+        @property
+        def _private(self):
+            return A()
+
+    pxy_a = make_final_proxy_type(
+        "A",
+        _Unusable,
+        A,
+        fast_to_slow=_Unusable(),
+        slow_to_fast=_Unusable(),
+    )
+
+    pxy_b = make_final_proxy_type(
+        "B",
+        _Unusable,
+        B,
+        fast_to_slow=_Unusable(),
+        slow_to_fast=_Unusable(),
+    )
+
+    result = pxy_b()._private
+    assert isinstance(result, pxy_a)
+
+
 def tuple_with_attrs(name, fields: list[str], extra_fields: set[str]):
     # Build a tuple-like class with some extra attributes and a custom
     # pickling scheme with __getnewargs_ex__