Decorators to save and load input of a function to and from a pickle file. Useful for debugging, prototyping or analysis

pickle_decorators

def save_input(fn):
    """saves input arguments as a pickle file for debugging/ analysis"""
    @wraps(fn)
    def wrapper(*args, **kwargs):

        # create tmp folder if it doesn't exist
        if not os.path.exists('tmp/pickles'):
            os.mkdir('tmp/pickles')

        # save input arguments
        with open(f'tmp/pickles/input_{fn.__name__}.pickle', 'wb') as f:
            pickle.dump((args, kwargs), f)
            
        # run function
        return fn(*args, **kwargs)
    return wrapper


def load_arguments_from_pickle(fn):
    """loads arguments from a pickle file for debugging/ analysis"""
    @wraps(fn)
    def wrapper(*args, **kwargs):
        
        # load input arguments
        with open(f'tmp/pickles/input_{fn.__name__}.pickle', 'rb') as f:
            args, kwargs = pickle.load(f)
            
        # run function
        return fn(*args, **kwargs)
    return wrapper


def test_pickle():
    @save_input
    def test(a, b, c):
        print('Running test and saving input as pickle')
        return a + b + c
    print(test(1, 2, 3))

    @load_arguments_from_pickle
    def test(a, b, c):
        print('Running test with loaded input from pickle')
        return a + b + c
    print(test())
    
if __name__ == '__main__':
    test_pickle()