Preserve descriptor binding in weak proxies for decorated callables

docs/changes.rst

@@ -107,6 +107,23 @@ Version 2.4.2
   implementation. The modulo argument of the three argument form of
   ``pow()`` is still not unwrapped, as described in the known issues.
 
+* Creating a ``WeakFunctionProxy`` around a decorated function, or a
+  decorated method, classmethod or staticmethod accessed via the class
+  rather than an instance, failed with ``TypeError``. In these cases the
+  function wrapper produced by the decorator has no bound instance, and the
+  proxy attempted to take a weak reference to ``None``. Only a decorated
+  method accessed via an instance worked.
+
+  The proxy now takes a weak reference to the instance only where there is
+  one, and also retains a weak reference to the class the wrapper was
+  accessed through. When called, the function is rebound against both, so
+  a decorated classmethod receives the class it was accessed through,
+  including a subclass, and a decorated instance method accessed via the
+  class recognises an instance passed as the first argument. The class is
+  not kept alive by the proxy, and if it is garbage collected a call raises
+  ``ReferenceError``, and the expiry callback runs, in the same way as when
+  the instance a method was bound to is garbage collected.
+
 Version 2.4.1
 -------------
 

src/wrapt/weakrefs.py

@@ -13,6 +13,14 @@
 # reference is therefore applied to the instance the method is bound to
 # and the original function. The function is then rebound at the point
 # of a call via the weak function proxy.
+#
+# Where the function is a wrapt function wrapper, such as results from
+# applying a decorator, the same applies but there may be no instance,
+# either because the wrapper was never bound, as for a decorated free
+# function, or because the method was accessed via the class rather
+# than an instance. In that case a weak reference to the class the
+# wrapper was accessed through, the owner, is retained as well so that
+# the function can still be rebound correctly at the point of a call.
 
 
 def _weak_function_proxy_callback(ref, proxy, callback):
@@ -62,9 +70,27 @@ def __init__(self, wrapped, callback=None):
         )
 
         self._self_expired = False
+        self._self_owner = None
 
         if isinstance(wrapped, _FunctionWrapperBase):
-            self._self_instance = weakref.ref(wrapped._self_instance, _callback)
+            # A function wrapper may have no instance, either because it
+            # was never bound, as for a decorated free function, or
+            # because the method was accessed via the class rather than
+            # an instance. Only take a weak reference to the instance
+            # where there is one. The owner is the class the wrapper was
+            # accessed through, and is retained so the function can be
+            # rebound with it when called. Without it a classmethod
+            # accessed via the class could not be rebound at all, and an
+            # instance method accessed via the class would not be able
+            # to identify an instance passed as the first argument.
+
+            instance = wrapped._self_instance
+            self._self_instance = (
+                weakref.ref(instance, _callback) if instance is not None else None
+            )
+            owner = wrapped._self_owner
+            if owner is not None:
+                self._self_owner = weakref.ref(owner, _callback)
 
             if wrapped._self_parent is not None:
                 # Explicit class in super() is used because the proxy
@@ -114,10 +140,24 @@ def _unpack_self(self, *args):
         if self._self_instance is not None and instance is None:
             raise ReferenceError("weakly-referenced object no longer exists")
 
-        # If the wrapped function was originally a bound function, for
-        # which we retained a reference to the instance and the unbound
-        # function we need to rebind the function and then call it. If
-        # not just called the wrapped function.
+        # If the wrapped function was a function wrapper for which the
+        # owner was retained, rebind the function against the instance,
+        # which may be None, and that owner. This is what a classmethod
+        # accessed via the class needs to be rebound at all, and what an
+        # instance method accessed via the class needs to recognise an
+        # instance passed as the first argument. If the owner has been
+        # garbage collected, raise a ReferenceError as for the instance.
+
+        if self._self_owner is not None:
+            owner = self._self_owner()
+            if owner is None:
+                raise ReferenceError("weakly-referenced object no longer exists")
+            return function.__get__(instance, owner)(*args, **kwargs)
+
+        # Otherwise, if the wrapped function was originally a bound
+        # function, for which we retained a reference to the instance and
+        # the unbound function, we need to rebind the function and then
+        # call it. If not just call the wrapped function.
 
         if instance is None:
             return self.__wrapped__(*args, **kwargs)

tests/core/test_weak_function_proxy.py

@@ -1,11 +1,95 @@
 import gc
+import sys
+import sysconfig
 import unittest
+import weakref
 
 import wrapt
 
 
 class TestWeakFunctionProxy(unittest.TestCase):
 
+    def test_decorated_function(self):
+        @wrapt.decorator
+        def decorator(wrapped, instance, args, kwargs):
+            return "decorated", wrapped(*args, **kwargs)
+
+        @decorator
+        def function(value):
+            return value
+
+        callbacks = []
+        proxy = wrapt.WeakFunctionProxy(function, lambda ref: callbacks.append(id(ref)))
+        self.assertEqual(proxy(42), ("decorated", 42))
+        del function
+        gc.collect()
+        self.assertEqual(callbacks, [id(proxy)])
+        with self.assertRaises(ReferenceError):
+            proxy(42)
+
+    def test_decorated_descriptors(self):
+        @wrapt.decorator
+        def decorator(wrapped, instance, args, kwargs):
+            return instance, wrapped(*args, **kwargs)
+
+        class Class:
+            @decorator
+            def method(self, value):
+                return value
+
+            @decorator
+            @classmethod
+            def class_method(cls, value):
+                return cls, value
+
+            @decorator
+            @staticmethod
+            def static_method(value):
+                return value
+
+        class Subclass(Class):
+            pass
+
+        obj = Class()
+        targets = ((Class, Class), (obj, Class), (Subclass, Subclass), (Subclass(), Subclass))
+        for target, owner in targets:
+            with self.subTest(target=target, kind="classmethod"):
+                proxy = wrapt.WeakFunctionProxy(target.class_method)
+                self.assertEqual(proxy(42), (owner, (owner, 42)))
+            with self.subTest(target=target, kind="staticmethod"):
+                proxy = wrapt.WeakFunctionProxy(target.static_method)
+                self.assertEqual(proxy(42), (None, 42))
+
+        proxy = wrapt.WeakFunctionProxy(Class.method)
+        self.assertEqual(proxy(obj, 42), (obj, 42))
+
+    @unittest.skipIf(
+        sys.version_info[:2] == (3, 13)
+        and sysconfig.get_config_var("Py_GIL_DISABLED"),
+        "Free-threaded CPython 3.13 immortalizes classes after a thread starts",
+    )
+    def test_decorated_classmethod_does_not_retain_owner(self):
+        @wrapt.decorator
+        def decorator(wrapped, instance, args, kwargs):
+            return wrapped(*args, **kwargs)
+
+        class Class:
+            @decorator
+            @classmethod
+            def method(cls):
+                return 42
+
+        callbacks = []
+        owner = weakref.ref(Class)
+        proxy = wrapt.WeakFunctionProxy(Class.method, lambda ref: callbacks.append(id(ref)))
+        self.assertEqual(proxy(), 42)
+        del Class
+        gc.collect()
+        self.assertIsNone(owner())
+        self.assertEqual(callbacks, [id(proxy)])
+        with self.assertRaises(ReferenceError):
+            proxy()
+
     def test_isinstance(self):
         def function(a, b):
             return a, b