Avoid holding strong reference to decorated and partial methods (#29)

* fix problem with names of slot

* ignore pycharm files

* partial unwrap

* add missed check

* fix annotation errors

* immprove type annotation

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Apply suggestions from code review

Co-authored-by: Talley Lambert <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix names

* sue __closure__

Co-authored-by: Talley Lambert <[email protected]>

* parametrize test

* ignore

* faster?

* catch assigned methods

* coverage

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Talley Lambert <[email protected]>

.gitignore

@@ -1,3 +1,4 @@
+.idea/
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

Makefile

@@ -35,4 +35,5 @@ benchmark-all:
 
 # compare HEAD against main
 benchmark-compare:
+	asv run -j 4 --interleave-processes --skip-existing main..HEAD --steps 2
 	asv compare --split --factor 1.1 main HEAD

psygnal/_signal.py

@@ -5,12 +5,14 @@
 import warnings
 import weakref
 from contextlib import contextmanager
-from functools import lru_cache, reduce
-from inspect import Parameter, Signature, ismethod
+from functools import lru_cache, partial, reduce
+from inspect import Parameter, Signature
+from types import MethodType
 from typing import (
     TYPE_CHECKING,
     Any,
     Callable,
+    Dict,
     Iterator,
     List,
     NoReturn,
@@ -24,14 +26,28 @@
 
 from typing_extensions import Literal
 
-MethodRef = Tuple["weakref.ReferenceType[object]", str]
+MethodRef = Tuple["weakref.ReferenceType[object]", Union[Callable, str]]
 NormedCallback = Union[MethodRef, Callable]
 StoredSlot = Tuple[NormedCallback, Optional[int]]
 AnyType = Type[Any]
 ReducerFunc = Callable[[tuple, tuple], tuple]
 _NULL = object()
 
 
+class PartialBoundMethodMeta(type):
+    def __instancecheck__(cls, inst: object) -> bool:
+        return isinstance(inst, partial) and isinstance(inst.func, MethodType)
+
+
+class PartialBoundMethod(metaclass=PartialBoundMethodMeta):
+    func: MethodType
+    args: List[Any]
+    keywords: Dict
+
+    def __call__(self, *args: List[Any], **kwargs: Dict[str, Any]) -> Any:
+        raise NotImplementedError  # pragma: no cover
+
+
 def signature(obj: Any) -> inspect.Signature:
     try:
         return inspect.signature(obj)
@@ -469,8 +485,10 @@ def _raise_connection_error(self, slot: Callable, extra: str = "") -> NoReturn:
         raise ValueError(msg)
 
     def _normalize_slot(self, slot: NormedCallback) -> NormedCallback:
-        if ismethod(slot):
-            return (weakref.ref(slot.__self__), slot.__name__)  # type: ignore
+        if isinstance(slot, MethodType):
+            return _get_proper_name(slot)
+        if isinstance(slot, PartialBoundMethod):
+            return partial_weakref(slot)
         if isinstance(slot, tuple) and not isinstance(slot[0], weakref.ref):
             return (weakref.ref(slot[0]), slot[1])
         return slot
@@ -648,22 +666,24 @@ def __call__(
         )
 
     def _run_emit_loop(self, args: Tuple[Any, ...]) -> None:
-
         rem: List[NormedCallback] = []
         # allow receiver to query sender with Signal.current_emitter()
         with self._lock:
             with Signal._emitting(self):
                 for (slot, max_args) in self._slots:
                     if isinstance(slot, tuple):
-                        _ref, method_name = slot
+                        _ref, method = slot
                         obj = _ref()
                         if obj is None:
                             rem.append(slot)  # add dead weakref
                             continue
-                        cb = getattr(obj, method_name, None)
-                        if cb is None:  # pragma: no cover
-                            rem.append(slot)  # object has changed?
-                            continue
+                        if callable(method):
+                            cb = method
+                        else:
+                            cb = getattr(obj, method, None)
+                            if cb is None:  # pragma: no cover
+                                rem.append(slot)  # object has changed?
+                                continue
                     else:
                         cb = slot
 
@@ -913,6 +933,39 @@ def _is_subclass(left: AnyType, right: type) -> bool:
     return issubclass(left, right)
 
 
+def partial_weakref(slot_partial: PartialBoundMethod) -> Tuple[weakref.ref, Callable]:
+    ref, name = _get_proper_name(slot_partial.func)
+    args_ = slot_partial.args
+    kwargs_ = slot_partial.keywords
+
+    def wrap(*args: Any, **kwargs: Any) -> Any:
+        getattr(ref(), name)(*args_, *args, **kwargs_, **kwargs)
+
+    return (ref, wrap)
+
+
+def _get_proper_name(slot: MethodType) -> Tuple[weakref.ref, str]:
+    obj = slot.__self__
+    # some decorators will alter method.__name__, so that obj.method
+    # will not be equal to getattr(obj, obj.method.__name__).
+    # We check for that case here and find the proper name in the function's closures
+    if getattr(obj, slot.__name__, None) != slot:
+        for c in slot.__closure__ or ():  # type: ignore
+            cname = getattr(c.cell_contents, "__name__", None)
+            if cname and getattr(obj, cname, None) == slot:
+                return weakref.ref(obj), cname
+        # slower, but catches cases like assigned functions
+        # that won't have function in closure
+        for name in reversed(dir(obj)):  # most dunder methods come first
+            if getattr(obj, name) == slot:
+                return weakref.ref(obj), name
+        # we don't know what to do here.
+        raise RuntimeError(  # pragma: no cover
+            f"Could not find method on {obj} corresponding to decorated function {slot}"
+        )
+    return weakref.ref(obj), slot.__name__
+
+
 try:
     import cython
 except ImportError:  # pragma: no cover

tests/test_psygnal.py

@@ -1,6 +1,7 @@
 import gc
 import time
 import weakref
+from functools import partial, wraps
 from inspect import Signature
 from types import FunctionType
 from typing import Optional
@@ -9,6 +10,24 @@
 import pytest
 
 from psygnal import Signal, SignalInstance
+from psygnal._signal import _get_proper_name
+
+
+def stupid_decorator(fun):
+    def _fun(*args):
+        fun(*args)
+
+    _fun.__annotations__ = fun.__annotations__
+    _fun.__name__ = "f_no_arg"
+    return _fun
+
+
+def good_decorator(fun):
+    @wraps(fun)
+    def _fun(*args):
+        fun(*args)
+
+    return _fun
 
 
 # fmt: off
@@ -36,6 +55,11 @@ def f_arg_kwarg(self, a, b=None): ...
     def f_vararg(self, *a): ...
     def f_vararg_varkwarg(self, *a, **b): ...
     def f_vararg_kwarg(self, *a, b=None): ...
+    @stupid_decorator
+    def f_int_decorated_stupid(self, a: int): ...
+    @good_decorator
+    def f_int_decorated_good(self, a: int): ...
+    f_any_assigned = lambda self, a: None  # noqa
 
 
 def f_no_arg(): ...
@@ -265,13 +289,27 @@ def test_signal_instance_error():
     assert "Signal() class attribute" in str(e)
 
 
-def test_weakrefs():
-    """Test that connect an instance method doesn't hold strong ref."""
+@pytest.mark.parametrize(
+    "slot",
+    [
+        "f_no_arg",
+        "f_int_decorated_stupid",
+        "f_int_decorated_good",
+        "f_any_assigned",
+        "partial",
+    ],
+)
+def test_weakref(slot):
+    """Test that a connected method doesn't hold strong ref."""
     emitter = Emitter()
     obj = MyObj()
 
     assert len(emitter.one_int) == 0
-    emitter.one_int.connect(obj.f_no_arg)
+    emitter.one_int.connect(
+        partial(obj.f_int_int, 1) if slot == "partial" else getattr(obj, slot)
+    )
+    assert len(emitter.one_int) == 1
+    emitter.one_int.emit(1)
     assert len(emitter.one_int) == 1
     del obj
     gc.collect()
@@ -571,3 +609,10 @@ def test_debug_import(monkeypatch):
     import psygnal
 
     assert not psygnal._compiled
+
+
+def test_get_proper_name():
+    obj = MyObj()
+    assert _get_proper_name(obj.f_int_decorated_stupid)[1] == "f_int_decorated_stupid"
+    assert _get_proper_name(obj.f_int_decorated_good)[1] == "f_int_decorated_good"
+    assert _get_proper_name(obj.f_any_assigned)[1] == "f_any_assigned"