Skip to content

Instantly share code, notes, and snippets.

@microprediction
Created September 24, 2020 20:38
Show Gist options
  • Save microprediction/48bcc3b9b66afb52d55f35e776ec3742 to your computer and use it in GitHub Desktop.
Save microprediction/48bcc3b9b66afb52d55f35e776ec3742 to your computer and use it in GitHub Desktop.
Decorator that caches intermediate research dataframes
CACHE = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))+'/cache/'
def picklecached():
"""
A function that creates a decorator which will use "cachefile" for caching the results of the decorated function "fn".
Also stores a CSV so we can quickly inspect the result ... though only pickles are used in loading.
"""
def decorator(fn): # define a decorator for a function "fn"
def wrapped(*args, **kwargs): # define a wrapper that will finally call "fn" with all arguments
# if cache exists -> load it and return its content
# Create a neame for the cache file
cachefile = CACHE + fn.__name__
if len(kwargs):
cachefile += '?'+'&'.join([k+'='+str(v) for k,v in kwargs.items()])
cachefile += '.pkl'
# Check if is already exists
if os.path.exists(cachefile):
with open(cachefile, 'rb') as cachehandle:
print("Using cached result from '%s'" % cachefile)
return pickle.load(cachehandle)
# execute the function with all arguments passed
res = fn(*args, **kwargs)
# write to cache file
with open(cachefile, 'wb') as cachehandle:
print("Saving result to cache '%s'" % cachefile)
pickle.dump(res, cachehandle)
try:
res.to_csv(cachefile.replace('.pkl','.csv'))
except:
print('Could not save CSV but pickle was dumped.')
return res
return wrapped
return decorator
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment