shelved.py

import hashlib
import pickle
import shelve

from decorator import decorator

SHELF: shelve.DbfilenameShelf


def _make_key(version, func, args, kwargs):
    to_hash = (version, func.__qualname__, tuple(args), tuple(kwargs.items()))
    buffer = pickle.dumps(to_hash)
    return hashlib.sha256(buffer).hexdigest()


@decorator
def shelved(func, version=None, *args, **kwargs):
    key = _make_key(version, func, args, kwargs)
    if key in SHELF:
        return SHELF[key]
    value = func(*args, **kwargs)
    SHELF[key] = value
    return value


@shelved(version=124)
def heavy_computation(x, y, z):
    import time
    print('slow')
    time.sleep(2.0)
    return x * 100 + y * 10 + z


def main():
    global SHELF
    with shelve.open('yolo') as SHELF:
        assert heavy_computation(1, 2, 3) == 123
        assert heavy_computation(1, 2, 3) == 123


if __name__ == '__main__':
    main()