Chromium Code Reviews| OLD | NEW |
|---|---|
| (Empty) | |
| 1 # Copyright 2015 The Chromium Authors. All rights reserved. | |
|
stgao
2016/09/21 23:03:24
As this file is not related to gitiles, I'd keep a
wrengr
2016/09/28 00:59:49
I have to move it. Leaving it in ./common causes a
| |
| 2 # Use of this source code is governed by a BSD-style license that can be | |
| 3 # found in the LICENSE file. | |
| 4 | |
| 5 """This module provides a decorator to cache the results of a function. | |
| 6 | |
| 7 Examples: | |
| 8 1. Decorate a function: | |
| 9 @cache_decorator.Cached() | |
| 10 def Test(a): | |
| 11 return a + a | |
| 12 | |
| 13 Test('a') | |
| 14 Test('a') # Returns the cached 'aa'. | |
| 15 | |
| 16 2. Decorate a method in a class: | |
| 17 class Downloader(object): | |
| 18 def __init__(self, url, retries): | |
| 19 self.url = url | |
| 20 self.retries = retries | |
| 21 | |
| 22 @property | |
| 23 def identifier(self): | |
| 24 return self.url | |
| 25 | |
| 26 @cache_decorator.Cached(): | |
| 27 def Download(self, path): | |
| 28 return urllib2.urlopen(self.url + '/' + path).read() | |
| 29 | |
| 30 d1 = Downloader('http://url', 4) | |
| 31 d1.Download('path') | |
| 32 | |
| 33 d2 = Downloader('http://url', 5) | |
| 34 d2.Download('path') # Returned the cached downloaded data. | |
| 35 """ | |
| 36 | |
| 37 import cStringIO | |
| 38 import functools | |
| 39 import hashlib | |
| 40 import inspect | |
| 41 import logging | |
| 42 import pickle | |
| 43 import zlib | |
| 44 | |
| 45 from google.appengine.api import memcache | |
| 46 | |
| 47 | |
| 48 class Cacher(object): | |
| 49 """An interface to cache and retrieve data. | |
| 50 | |
| 51 Subclasses should implement the Get/Set functions. | |
| 52 TODO: Add a Delete function (default to no-op) if needed later. | |
| 53 """ | |
| 54 def Get(self, key): | |
| 55 """Returns the cached data for the given key if available. | |
| 56 | |
| 57 Args: | |
| 58 key (str): The key to identify the cached data. | |
| 59 """ | |
| 60 raise NotImplementedError() | |
| 61 | |
| 62 def Set(self, key, data, expire_time=0): | |
| 63 """Cache the given data which is identified by the given key. | |
| 64 | |
| 65 Args: | |
| 66 key (str): The key to identify the cached data. | |
| 67 data (object): The python object to be cached. | |
| 68 expire_time (int): Number of seconds from current time (up to 1 month). | |
| 69 """ | |
| 70 raise NotImplementedError() | |
| 71 | |
| 72 | |
| 73 class PickledMemCacher(Cacher): | |
| 74 """A memcache-backed implementation of the interface Cacher. | |
| 75 | |
| 76 The data to be cached should be pickleable. | |
| 77 Limitation: size of the pickled data and key should be <= 1MB. | |
| 78 """ | |
| 79 def Get(self, key): | |
| 80 return memcache.get(key) | |
| 81 | |
| 82 def Set(self, key, data, expire_time=0): | |
| 83 return memcache.set(key, data, time=expire_time) | |
| 84 | |
| 85 | |
| 86 class _CachedItemMetaData(object): | |
| 87 def __init__(self, number): | |
| 88 self.number = number | |
| 89 | |
| 90 | |
| 91 class CompressedMemCacher(Cacher): | |
| 92 """A memcache-backed implementation of the interface Cacher with compression. | |
| 93 | |
| 94 The data to be cached would be pickled and then compressed. | |
| 95 Data still > 1MB will be split into sub-piece and stored separately. | |
| 96 During retrieval, if any sub-piece is missing, None is returned. | |
| 97 """ | |
| 98 CHUNK_SIZE = 990000 | |
| 99 | |
| 100 def Get(self, key): | |
| 101 data = memcache.get(key) | |
| 102 if isinstance(data, _CachedItemMetaData): | |
| 103 num = data.number | |
| 104 sub_keys = ['%s-%s' % (key, i) for i in range(num)] | |
| 105 all_data = memcache.get_multi(sub_keys) | |
| 106 if len(all_data) != num: # Some data is missing. | |
| 107 return None | |
| 108 | |
| 109 data_output = cStringIO.StringIO() | |
| 110 for sub_key in sub_keys: | |
| 111 data_output.write(all_data[sub_key]) | |
| 112 data = data_output.getvalue() | |
| 113 | |
| 114 return None if data is None else pickle.loads(zlib.decompress(data)) | |
| 115 | |
| 116 def Set(self, key, data, expire_time=0): | |
| 117 pickled_data = pickle.dumps(data) | |
| 118 compressed_data = zlib.compress(pickled_data) | |
| 119 | |
| 120 all_data = {} | |
| 121 if len(compressed_data) > self.CHUNK_SIZE: | |
| 122 num = 0 | |
| 123 for index in range(0, len(compressed_data), self.CHUNK_SIZE): | |
| 124 sub_key = '%s-%s' % (key, num) | |
| 125 all_data[sub_key] = compressed_data[index : index + self.CHUNK_SIZE] | |
| 126 num += 1 | |
| 127 | |
| 128 all_data[key] = _CachedItemMetaData(num) | |
| 129 else: | |
| 130 all_data[key] = compressed_data | |
| 131 | |
| 132 keys_not_set = memcache.set_multi(all_data, time=expire_time) | |
| 133 return len(keys_not_set) == 0 | |
| 134 | |
| 135 | |
| 136 def _DefaultKeyGenerator(func, args, kwargs): | |
| 137 """Generates a key from the function and arguments passed to it. | |
| 138 | |
| 139 Args: | |
| 140 func (function): An arbitrary function. | |
| 141 args (list): Positional arguments passed to ``func``. | |
| 142 kwargs (dict): Keyword arguments passed to ``func``. | |
| 143 | |
| 144 Returns: | |
| 145 A string to represent a call to the given function with the given arguments. | |
| 146 """ | |
| 147 params = inspect.getcallargs(func, *args, **kwargs) | |
| 148 for var_name in params: | |
| 149 if not hasattr(params[var_name], 'identifier'): | |
| 150 continue | |
| 151 | |
| 152 if callable(params[var_name].identifier): | |
| 153 params[var_name] = params[var_name].identifier() | |
| 154 else: | |
| 155 params[var_name] = params[var_name].identifier | |
| 156 | |
| 157 return hashlib.md5(pickle.dumps(params)).hexdigest() | |
| 158 | |
| 159 | |
| 160 def Cached(namespace=None, | |
| 161 expire_time=0, | |
| 162 key_generator=_DefaultKeyGenerator, | |
| 163 cacher=PickledMemCacher()): | |
| 164 """Returns a decorator to cache the decorated function's results. | |
| 165 | |
| 166 However, if the function returns None, empty list/dict, empty string, or other | |
| 167 value that is evaluated as False, the results won't be cached. | |
| 168 | |
| 169 This decorator is to cache results of different calls to the decorated | |
| 170 function, and avoid executing it again if the calls are equivalent. Two calls | |
| 171 are equivalent, if the namespace is the same and the keys generated by the | |
| 172 ``key_generator`` are the same. | |
| 173 | |
| 174 The usage of this decorator requires that: | |
| 175 - If the default key generator is used, parameters passed to the decorated | |
| 176 function should be pickleable, or each of the parameter has an identifier | |
| 177 property or method which returns pickleable results. | |
| 178 - If the default cacher is used, the returned results of the decorated | |
| 179 function should be pickleable. | |
| 180 | |
| 181 Args: | |
| 182 namespace (str): A prefix to the key for the cache. Default to the | |
| 183 combination of module name and function name of the decorated function. | |
| 184 expire_time (int): Expiration time, relative number of seconds from current | |
| 185 time (up to 1 month). Defaults to 0 -- never expire. | |
| 186 key_generator (function): A function to generate a key to represent a call | |
| 187 to the decorated function. Defaults to :func:`_DefaultKeyGenerator`. | |
| 188 cacher (Cacher): An instance of an implementation of interface `Cacher`. | |
| 189 Defaults to one of `PickledMemCacher` which is based on memcache. | |
| 190 | |
| 191 Returns: | |
| 192 The cached results or the results of a new run of the decorated function. | |
| 193 """ | |
| 194 def GetPrefix(func, namespace): | |
| 195 return namespace or '%s.%s' % (func.__module__, func.__name__) | |
| 196 | |
| 197 def Decorator(func): | |
| 198 """Decorator to cache a function's results.""" | |
| 199 @functools.wraps(func) | |
| 200 def Wrapped(*args, **kwargs): | |
| 201 prefix = GetPrefix(func, namespace) | |
| 202 key = '%s-%s' % (prefix, key_generator(func, args, kwargs)) | |
| 203 | |
| 204 try: | |
| 205 result = cacher.Get(key) | |
| 206 except Exception: # pragma: no cover. | |
| 207 result = None | |
| 208 logging.exception( | |
| 209 'Failed to get cached data for function %s.%s, args=%s, kwargs=%s', | |
| 210 func.__module__, func.__name__, repr(args), repr(kwargs)) | |
| 211 | |
| 212 if result is not None: | |
| 213 return result | |
| 214 | |
| 215 result = func(*args, **kwargs) | |
| 216 if result: | |
| 217 try: | |
| 218 cacher.Set(key, result, expire_time=expire_time) | |
| 219 except Exception: # pragma: no cover. | |
| 220 logging.exception( | |
| 221 'Failed to cache data for function %s.%s, args=%s, kwargs=%s', | |
| 222 func.__module__, func.__name__, repr(args), repr(kwargs)) | |
| 223 | |
| 224 return result | |
| 225 | |
| 226 return Wrapped | |
| 227 | |
| 228 return Decorator | |
| OLD | NEW |