e708fc44dc3ebdfec4f1eceb7334ec02c55b464b
[leap_pycommon.git] / 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 functools
22 import logging
23
24 logger = logging.getLogger(__name__)
25
26
27 class _memoized(object):
28     """
29     Decorator.
30
31     Caches a function's return value each time it is called.
32     If called later with the same arguments, the cached value is returned
33     (not reevaluated).
34     """
35     def __init__(self, func, ignore_kwargs=None, is_method=False):
36         """
37         :param ignore_kwargs: If True, ignore all kwargs.
38                               If tuple, ignore those kwargs.
39         :type ignore_kwargs: bool, tuple or None
40         :param is_method: whether the decorated function is a method.
41                           (ignores the self argument if so).
42         :type is_method: True
43         """
44         self.ignore_kwargs = ignore_kwargs if ignore_kwargs else []
45         self.is_method = is_method
46         self.func = func
47
48         # TODO should put bounds to the cache dict so we do not
49         # consume a huge amount of memory.
50         self.cache = {}
51
52     def __call__(self, *args, **kwargs):
53         """
54         Executes the call.
55
56         :tyoe args: tuple
57         :type kwargs: dict
58         """
59         def ret_or_raise(value):
60             """
61             Returns the value except if it is an exception,
62             in which case it's raised.
63             """
64             if isinstance(value, Exception):
65                 raise value
66             return value
67
68         if self.is_method:
69             # forget about `self` as key
70             key_args = args[1:]
71         if self.ignore_kwargs is True:
72             key = key_args
73         else:
74             key = (key_args, frozenset(
75                 [(k, v) for k, v in kwargs.items()
76                  if k not in self.ignore_kwargs]))
77
78         if not isinstance(key, collections.Hashable):
79             # uncacheable. a list, for instance.
80             # better to not cache than blow up.
81             logger.warning("Key is not hashable, bailing out!")
82             return self.func(*args, **kwargs)
83
84         if key in self.cache:
85             logger.debug("Got value from cache...")
86             value = self.cache[key]
87             return ret_or_raise(value)
88         else:
89             try:
90                 value = self.func(*args, **kwargs)
91             except Exception as exc:
92                 logger.error("Exception while calling function: %r" % (exc,))
93                 value = exc
94             self.cache[key] = value
95             return ret_or_raise(value)
96
97     def __repr__(self):
98         """
99         Return the function's docstring.
100         """
101         return self.func.__doc__
102
103     def __get__(self, obj, objtype):
104         """
105         Support instance methods.
106         """
107         return functools.partial(self.__call__, obj)
108
109
110 def memoized_method(function=None, ignore_kwargs=None):
111     """
112     Wrap _memoized to allow for deferred calling
113
114     :type function: callable, or None.
115     :type ignore_kwargs: None, True or tuple.
116     """
117     if function:
118         return _memoized(function, is_method=True)
119     else:
120         def wrapper(function):
121             return _memoized(
122                 function, ignore_kwargs=ignore_kwargs, is_method=True)
123         return wrapper