Improve and document decorator

Author: edan-bainglass

src/aiidalab_qe/common/decorators.py

@@ -14,32 +14,70 @@ def cache_per_thread(invalidator: str | None = None):
     `invalidator` : `str | None`
         The name of the attribute to watch for changes.
         If the attribute changes, the cache will be invalidated.
+
+    Usage
+    -----
+    Methods / free functions:
+    >>> @cache_per_thread
+    >>> def f(...): ...
+
+    Properties (must use this order):
+    >>> @cache_per_thread(invalidator="uuid")
+    >>> @property
+    >>> def f(...): ...
     """
 
-    def decorator(func):
-        @functools.wraps(func)
-        def wrapper(self, *args, **kwargs):
-            cache = getattr(_thread_local, "cache", None)
-            if cache is None:
-                cache = {}
-                _thread_local.cache = cache
+    def cached_call(func, *args, **kwargs):
+        # Get or initialize the cache for this thread
+        cache = getattr(_thread_local, "cache", None)
+        if cache is None:
+            cache = {}
+            _thread_local.cache = cache
+
+        # Begin constructing the cache key
+        key_parts = [func]
+
+        if args:
+            # Check if func is a method of a class
+            if hasattr(args[0], "__dict__"):
+                self = args[0]
+                key_parts.append(id(self))
+
+                # Include the invalidator attribute if specified
+                if invalidator is not None:
+                    key_parts.append(getattr(self, invalidator))
 
-            key_parts = [id(self), func]
-            if invalidator is not None:
-                key_parts.append(getattr(self, invalidator))
-            if args or kwargs:
-                key_parts.append(args)
-                key_parts.append(frozenset(kwargs.items()))
-            key = tuple(key_parts)
+        # Include the arguments in the cache key
+        if args or kwargs:
+            key_parts.append(args)
+            key_parts.append(frozenset(kwargs.items()))
 
-            if key not in cache:
-                cache[key] = func(self, *args, **kwargs)
+        key = tuple(key_parts)
 
-            return cache[key]
+        # Check if the result is already cached
+        if key not in cache:
+            cache[key] = func(*args, **kwargs)
 
+        return cache[key]
+
+    def decorator(func):
+        # Property case
         if isinstance(func, property):
+            fget = func.fget
+            if fget is None:
+                raise TypeError("Property has no getter to wrap")
+
+            @functools.wraps(fget)
+            def wrapper(self):  # type: ignore
+                return cached_call(fget, self)
+
             return property(wrapper)
 
+        # method/function case
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            return cached_call(func, *args, **kwargs)
+
         return wrapper
 
     return decorator

tests/test_decorators.py

@@ -8,29 +8,40 @@ def __init__(self, value):
         self.value = value
 
     @cache_per_thread("value")
-    def compute(self, x):
+    @property
+    def double_the_value(self):
+        return self.value * 2
+
+    @cache_per_thread("value")
+    def add_to_value(self, x):
         return self.value + x
 
 
+def test_cache_per_thread_property():
+    obj = DummyClass(10)
+    assert obj.double_the_value == 20  # First call, computes the result
+    assert obj.double_the_value == 20  # Cached result, no recomputation
+
+
 def test_cache_per_thread_single_thread():
     obj = DummyClass(10)
-    assert obj.compute(5) == 15  # First call, computes the result
-    assert obj.compute(5) == 15  # Cached result, no recomputation
+    assert obj.add_to_value(5) == 15  # First call, computes the result
+    assert obj.add_to_value(5) == 15  # Cached result, no recomputation
 
 
 def test_cache_per_thread_different_args():
     obj = DummyClass(10)
-    assert obj.compute(5) == 15  # First call with x=5
-    assert obj.compute(3) == 13  # First call with x=3
-    assert obj.compute(5) == 15  # Cached result for x=5
+    assert obj.add_to_value(5) == 15  # First call with x=5
+    assert obj.add_to_value(3) == 13  # First call with x=3
+    assert obj.add_to_value(5) == 15  # Cached result for x=5
 
 
 def test_cache_per_thread_different_threads():
     obj = DummyClass(300)  # must be > 256 to avoid Python integer interning
     results = []
 
     def thread_func(x):
-        results.append(obj.compute(x))
+        results.append(obj.add_to_value(x))
 
     thread1 = Thread(target=thread_func, args=(5,))
     thread2 = Thread(target=thread_func, args=(5,))
@@ -50,7 +61,16 @@ def thread_func(x):
 
 def test_cache_invalidation():
     obj = DummyClass(10)
-    assert obj.compute(5) == 15  # First call, computes the result
-    assert obj.compute(5) == 15  # Cached result, no recomputation
+    assert obj.add_to_value(5) == 15  # First call, computes the result
+    assert obj.add_to_value(5) == 15  # Cached result, no recomputation
     obj.value = 20  # Invalidate cache by changing the invalidator attribute
-    assert obj.compute(5) == 25  # Recomputes the result
+    assert obj.add_to_value(5) == 25  # Recomputes the result
+
+
+def test_cache_per_thread_on_plain_function():
+    @cache_per_thread()
+    def compute(x):
+        return 42 + x
+
+    assert compute(5) == 47  # First call, computes the result
+    assert compute(5) == 47  # Cached result, no recomputation