Add support for deferreds to memoize_method decorator

2 files changed, 67 insertions, 26 deletions

diff --git a/changes/feature-memoize_deferred b/changes/feature-memoize_deferred
new file mode 100644
index 0000000..dd2fdee
--- /dev/null
+++ b/changes/feature-memoize_deferred
@@ -0,0 +1 @@
+- Add support for deferreds to memoize_method decorator
diff --git a/src/leap/common/decorators.py b/src/leap/common/decorators.py
index 99c3653..4b07ea9 100644
--- a/src/leap/common/decorators.py
+++ b/src/leap/common/decorators.py
@@ -22,6 +22,13 @@ import datetime
 import functools
 import logging
 
+try:
+    from twisted.internet import defer
+except ImportError:
+    class defer:
+        class Deferred:
+            pass
+
 logger = logging.getLogger(__name__)
 
 
@@ -59,6 +66,7 @@ class _memoized(object):
         # consume a huge amount of memory.
         self.cache = {}
         self.cache_ts = {}
+        self.is_deferred = {}
 
     def __call__(self, *args, **kwargs):
         """
@@ -67,28 +75,7 @@ class _memoized(object):
         :tyoe args: tuple
         :type kwargs: dict
         """
-        def ret_or_raise(value):
-            """
-            Returns the value except if it is an exception,
-            in which case it's raised.
-            """
-            if isinstance(value, Exception):
-                raise value
-            return value
-
-        if self.is_method:
-            # forget about `self` as key
-            key_args = args[1:]
-        else:
-            key_args = args
-
-        if self.ignore_kwargs is True:
-            key = key_args
-        else:
-            key = (key_args, frozenset(
-                [(k, v) for k, v in kwargs.items()
-                 if k not in self.ignore_kwargs]))
-
+        key = self._build_key(*args, **kwargs)
         if not isinstance(key, collections.Hashable):
             # uncacheable. a list, for instance.
             # better to not cache than blow up.
@@ -97,9 +84,9 @@ class _memoized(object):
 
         if key in self.cache:
             if self._is_cache_still_valid(key):
-                value = self.cache[key]
                 logger.debug("Got value from cache...")
-                return ret_or_raise(value)
+                value = self._get_value(key)
+                return self._ret_or_raise(value)
             else:
                 logger.debug("Cache is invalid, evaluating again...")
 
@@ -109,9 +96,52 @@ class _memoized(object):
         except Exception as exc:
             logger.error("Exception while calling function: %r" % (exc,))
             value = exc
-        self.cache[key] = value
+
+        if isinstance(value, defer.Deferred):
+            value.addBoth(self._store_deferred_value(key))
+        else:
+            self.cache[key] = value
+            self.is_deferred[key] = False
         self.cache_ts[key] = datetime.datetime.now()
-        return ret_or_raise(value)
+        return self._ret_or_raise(value)
+
+    def _build_key(self, *args, **kwargs):
+        """
+        Build key from the function arguments
+        """
+        if self.is_method:
+            # forget about `self` as key
+            key_args = args[1:]
+        else:
+            key_args = args
+
+        if self.ignore_kwargs is True:
+            key = key_args
+        else:
+            key = (key_args, frozenset(
+                [(k, v) for k, v in kwargs.items()
+                 if k not in self.ignore_kwargs]))
+        return key
+
+    def _get_value(self, key):
+        """
+        Get a value form cache for a key
+        """
+        if self.is_deferred[key]:
+            value = self.cache[key]
+            # it produces an errback if value is Failure
+            return defer.succeed(value)
+        else:
+            return self.cache[key]
+
+    def _ret_or_raise(self, value):
+        """
+        Returns the value except if it is an exception,
+        in which case it's raised.
+        """
+        if isinstance(value, Exception):
+            raise value
+        return value
 
     def _is_cache_still_valid(self, key, now=datetime.datetime.now):
         """
@@ -131,6 +161,16 @@ class _memoized(object):
         delta = datetime.timedelta(seconds=self.CACHE_INVALIDATION_DELTA)
         return (now() - cached_ts) < delta
 
+    def _store_deferred_value(self, key):
+        """
+        Returns a callback to store values from deferreds
+        """
+        def callback(value):
+            self.cache[key] = value
+            self.is_deferred[key] = True
+            return value
+        return callback
+
     def __repr__(self):
         """
         Return the function's docstring.