src/leap/common/decorators.py

   1 # -*- coding: utf-8 -*-
   2 # decorators.py
   3 # Copyright (C) 2013 LEAP
   4 #
   5 # This program is free software: you can redistribute it and/or modify
   6 # it under the terms of the GNU General Public License as published by
   7 # the Free Software Foundation, either version 3 of the License, or
   8 # (at your option) any later version.
   9 #
  10 # This program is distributed in the hope that it will be useful,
  11 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13 # GNU General Public License for more details.
  14 #
  15 # You should have received a copy of the GNU General Public License
  16 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
  17 """
  18 Useful decorators.
  19 """
  20 import collections
  21 import datetime
  22 import functools
  23 import logging
  24
  25 try:
  26     from twisted.internet import defer
  27 except ImportError:
  28     class defer:
  29         class Deferred:
  30             pass
  31
  32 logger = logging.getLogger(__name__)
  33
  34
  35 class _memoized(object):
  36     """
  37     Decorator.
  38
  39     Caches a function's return value each time it is called.
  40     If called later with the same arguments, the cached value is returned
  41     (not reevaluated).
  42     """
  43
  44     # cache invalidation time, in seconds
  45     CACHE_INVALIDATION_DELTA = 1800
  46
  47     def __init__(self, func, ignore_kwargs=None, is_method=False,
  48                  invalidation=None):
  49
  50         """
  51         :param ignore_kwargs: If True, ignore all kwargs.
  52                               If tuple, ignore those kwargs.
  53         :type ignore_kwargs: bool, tuple or None
  54         :param is_method: whether the decorated function is a method.
  55                           (ignores the self argument if so).
  56         :type is_method: True
  57         """
  58         self.ignore_kwargs = ignore_kwargs if ignore_kwargs else []
  59         self.is_method = is_method
  60         self.func = func
  61
  62         if invalidation:
  63             self.CACHE_INVALIDATION_DELTA = invalidation
  64
  65         # TODO should put bounds to the cache dict so we do not
  66         # consume a huge amount of memory.
  67         self.cache = {}
  68         self.cache_ts = {}
  69         self.is_deferred = {}
  70
  71     def __call__(self, *args, **kwargs):
  72         """
  73         Executes the call.
  74
  75         :tyoe args: tuple
  76         :type kwargs: dict
  77         """
  78         key = self._build_key(*args, **kwargs)
  79         if not isinstance(key, collections.Hashable):
  80             # uncacheable. a list, for instance.
  81             # better to not cache than blow up.
  82             logger.warning("Key is not hashable, bailing out!")
  83             return self.func(*args, **kwargs)
  84
  85         if key in self.cache:
  86             if self._is_cache_still_valid(key):
  87                 logger.debug("Got value from cache...")
  88                 value = self._get_value(key)
  89                 return self._ret_or_raise(value)
  90             else:
  91                 logger.debug("Cache is invalid, evaluating again...")
  92
  93         # no cache, or cache invalid
  94         try:
  95             value = self.func(*args, **kwargs)
  96         except Exception as exc:
  97             logger.error("Exception while calling function: %r" % (exc,))
  98             value = exc
  99
 100         if isinstance(value, defer.Deferred):
 101             value.addBoth(self._store_deferred_value(key))
 102         else:
 103             self.cache[key] = value
 104             self.is_deferred[key] = False
 105         self.cache_ts[key] = datetime.datetime.now()
 106         return self._ret_or_raise(value)
 107
 108     def _build_key(self, *args, **kwargs):
 109         """
 110         Build key from the function arguments
 111         """
 112         if self.is_method:
 113             # forget about `self` as key
 114             key_args = args[1:]
 115         else:
 116             key_args = args
 117
 118         if self.ignore_kwargs is True:
 119             key = key_args
 120         else:
 121             key = (key_args, frozenset(
 122                 [(k, v) for k, v in kwargs.items()
 123                  if k not in self.ignore_kwargs]))
 124         return key
 125
 126     def _get_value(self, key):
 127         """
 128         Get a value form cache for a key
 129         """
 130         if self.is_deferred[key]:
 131             value = self.cache[key]
 132             # it produces an errback if value is Failure
 133             return defer.succeed(value)
 134         else:
 135             return self.cache[key]
 136
 137     def _ret_or_raise(self, value):
 138         """
 139         Returns the value except if it is an exception,
 140         in which case it's raised.
 141         """
 142         if isinstance(value, Exception):
 143             raise value
 144         return value
 145
 146     def _is_cache_still_valid(self, key, now=datetime.datetime.now):
 147         """
 148         Returns True if the cache value is still valid, False otherwise.
 149
 150         For now, this happen if less than CACHE_INVALIDATION_DELTA seconds
 151         have passed from the time in which we recorded the cached value.
 152
 153         :param key: the key to lookup in the cache
 154         :type key: hashable
 155         :param now: a callable that returns a datetime object. override
 156                     for dependency injection during testing.
 157         :type now: callable
 158         :rtype: bool
 159         """
 160         cached_ts = self.cache_ts[key]
 161         delta = datetime.timedelta(seconds=self.CACHE_INVALIDATION_DELTA)
 162         return (now() - cached_ts) < delta
 163
 164     def _store_deferred_value(self, key):
 165         """
 166         Returns a callback to store values from deferreds
 167         """
 168         def callback(value):
 169             self.cache[key] = value
 170             self.is_deferred[key] = True
 171             return value
 172         return callback
 173
 174     def __repr__(self):
 175         """
 176         Return the function's docstring.
 177         """
 178         return self.func.__doc__
 179
 180     def __get__(self, obj, objtype):
 181         """
 182         Support instance methods.
 183         """
 184         return functools.partial(self.__call__, obj)
 185
 186
 187 def memoized_method(function=None, ignore_kwargs=None,
 188                     invalidation=_memoized.CACHE_INVALIDATION_DELTA):
 189     """
 190     Wrap _memoized to allow for deferred calling
 191
 192     :type function: callable, or None.
 193     :type ignore_kwargs: None, True or tuple.
 194     :type invalidation: int seconds.
 195     """
 196     if function:
 197         return _memoized(function, is_method=True, invalidation=invalidation)
 198     else:
 199         def wrapper(function):
 200             return _memoized(
 201                 function, ignore_kwargs=ignore_kwargs, is_method=True,
 202                 invalidation=invalidation)
 203         return wrapper