Source code for nnsa.utils.pkl

import pickle


def pickle_load(filepath):
    with open(filepath, 'rb') as f:
        a = pickle.load(f)
    return a


def pickle_save(filepath, a, overwrite=False, protocol='default', verbose=0):
    from nnsa.utils.paths import check_filename_exists, check_directory_exists

    if not overwrite:
        # Raise error if exists.
        check_filename_exists(filepath)
    check_directory_exists(filepath=filepath)

    # Select protocol.
    if isinstance(protocol, str):
        if protocol.lower() == 'default':
            protocol = pickle.DEFAULT_PROTOCOL
        elif protocol.lower() == 'highest':
            protocol = pickle.HIGHEST_PROTOCOL
        else:
            raise ValueError(f'Invalid value `protocol` = "{protocol}".')

    if verbose:
        print('Saving to pickle...')

    with open(filepath, 'wb') as f:
        pickle.dump(a, f, protocol=protocol)

    if verbose:
        print('Saved to {}!'.format(filepath))