Memoization to disk

memoize.py

import os
import gzip
import pickle
import dbm
import hashlib

MYDIR = "."  # dir to save files

def memoize(dbmfile=os.path.join(MYDIR, "memoize.db"), compress=True):
    "help save time and bandwidth on heavy functions"
    def _fn(fn):
        def _deco(*args, **kwargs):
            # Picle the function arguments in protocol 5 and obtain hash as the
            # store keys
            key = fn.__name__ + hashlib.md5(pickle.dumps((args, kwargs), 5)).hexdigest()
            # Open database and check if key exists
            if os.path.isfile(dbmfile):
                with dbm.open(dbmfile, "r") as db:
                    if blob := db.get(key):
                        return pickle.loads(gzip.decompress(blob) if compress else blob)
            # Call function and save into the store
            ret = fn(*args, **kwargs)
            with dbm.open(dbmfile, "c") as db:
                blob = pickle.dumps(ret)
                db[key] = gzip.compress(blob) if compress else blob
            return ret
        return _deco
    return _fn

# Example of use
import pandas_datareader as pdr

@memoize()
def ohlcv(ticker, startdate, enddate) -> pd.DataFrame:
    """Pull OHLCV from Yahoo Finance"""
    df = pdr.get_data_yahoo(ticker, start=startdate, end=enddate)
    df['symbol'] = ticker
    return df