Last active
December 18, 2019 11:09
-
-
Save raven4752/73654777dad7670f77471aa7b47695b2 to your computer and use it in GitHub Desktop.
python decorator function to cache function's returned value to disk
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
def cache_result(cache_file, param_in_suffix=None, root_dir=None): | |
""" | |
decorator function to cache function's return value (assume returned value can be pickled) | |
the cached file with be located to root_dir/[_param-key_param-value,]/cache_file.pkl | |
usage: | |
@cache_result(cache_file_name,param_in_suffix=[param_key],) | |
def function_to_create_cache(): | |
.... | |
:param root_dir: dir to store the cache if the wrapped function has root_dir as a param, the value will be used instead | |
:param cache_file: prefix of the cached file name | |
:param param_in_suffix: a list of params of the wrapped function to be suffixed to the cached file | |
:return: a wrapped function | |
""" | |
def _cache_result(func): | |
argspec = inspect.getfullargspec(func) | |
positional_count = len(argspec.args) - len(argspec.defaults) | |
defaults = dict(zip(argspec.args[positional_count:], argspec.defaults)) | |
def wrapped_func(*param, **kwargs): | |
if 'root_dir' in kwargs: | |
root_dir_ = kwargs['root_dir'] | |
else: | |
root_dir_ = root_dir | |
if param_in_suffix is not None: | |
suffix = '' | |
for k in param_in_suffix: | |
if k not in kwargs: | |
suffix_ = str(defaults[k]) | |
else: | |
suffix_ = str(kwargs[k]) | |
suffix = suffix + '_' + k + '_' + suffix_ | |
cache_path = join(root_dir_, suffix, cache_file) + '.pkl' | |
else: | |
cache_path = join(root_dir_, cache_file) + '.pkl' | |
if not exists(cache_path): | |
result = func(*param, **kwargs) | |
dir = dirname(cache_path) | |
if not exists(dir): | |
os.makedirs(dir) | |
with open(cache_path, 'wb') as f: | |
print('saving cache to {} when calling {}'.format(cache_path, | |
func.__module__ + '.' + func.__name__)) | |
pickle.dump(result, f) | |
return result | |
else: | |
print('loading cache from {} when calling {}'.format(cache_path, func.__module__ + '.' + func.__name__)) | |
with open(cache_path, 'rb') as f: | |
return pickle.load(f) | |
return wrapped_func | |
return _cache_result |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment