pickle_cache.py

"""
See __test_time function for example use

>>> import pathlib
>>> pathlib.Path("__test_time.pickle").unlink(missing_ok=True)
>>> a = __test_time()     # Get current time
>>> import time
>>> time.sleep(0.1)
>>> __test_time() - a     # Time should be the same since it is cached.
0.0
>>> import os
>>> round(time.time() - os.stat(".__main__.__test_time.pickle").st_mtime, 1)
0.1
>>> b = __test_time_b()
>>> time.sleep(0.2)
>>> round(time.time() - os.stat("__something_else.pickle").st_mtime, 1)
0.2
>>> pathlib.Path("__something_else.pickle").unlink(missing_ok=True)
>>> pathlib.Path("__test_time.pickle").unlink(missing_ok=True)
"""

import functools
import hashlib
import os
import pickle
import time
from datetime import timedelta

def pickle_cache(fn=None, *, cache_path=None, ttl=timedelta(hours=6)):
    if fn is None:
        return functools.partial(pickle_cache, cache_path=cache_path, ttl=ttl)

    @functools.wraps(fn)
    def wrapped(*args, **kwargs):
        nonlocal cache_path
        if cache_path is None:
            # Generate a unique cache filename based on package, module, function name, and arguments
            module_name = fn.__module__
            function_name = fn.__qualname__
            parts = []
            if args:
                parts.append(args)
            if kwargs:
                parts.append(kwargs)
            hash_args = '_' + hashlib.md5(str(parts).encode()).hexdigest() if parts else ''
            cache_path = f".{module_name}.{function_name}{hash_args}.pickle"

        if os.path.exists(cache_path):
            if time.time() - os.stat(cache_path).st_mtime < ttl.total_seconds():
                with open(cache_path, 'rb') as f:
                    return pickle.load(f)

        result = fn(*args, **kwargs)
        with open(cache_path, 'wb') as f:
            pickle.dump(result, f)
        return result
    return wrapped

# Example usage of the updated decorator
@pickle_cache
def __test_time():
    import time
    return time.time()

@pickle_cache(cache_path='__something_else.pickle')
def __test_time_b():
    import time
    return time.time()

if __name__ == "__main__":
    import doctest
    doctest.testmod()