persist_to_file.py

import os
import pickle
import time
import glob
import joblib
import inspect
from loguru import logger


def persist_to_file():
    def decorator(original_func):
        path = original_func.__name__

        os.makedirs(path, exist_ok=True)
        fns = [fn.split('/')[-1] for fn in glob.glob(f"{path}/*")]

        def new_func(*args, **kwargs):
            t0 = time.time()
            code = inspect.getsource(original_func)
            param_hash = joblib.hash(dict(args=args, kwargs=kwargs, code=code))
            t1 = time.time()
            dt = t1 - t0
            logger.debug(f"Took {dt:1.3f}s to hash args for {path}")
            if param_hash not in fns:
                ret = original_func(*args, **kwargs)
                t0 = time.time()
                pickle.dump(ret, open(path + '/' + param_hash, 'wb'))
                t1 = time.time()
                dt = t1 - t0
                logger.debug(f"Took {dt:1.3f}s to save args")
                return ret
            else:
                ret = pickle.load(open(path + '/' + param_hash, 'rb'))
                return ret
        return new_func
    return decorator