Source code for secml.utils.pickle_utils

"""
.. module:: PickleWrapper
   :synopsis: Wrapper for cPickle object saving package

.. moduleauthor:: Ambra Demontis <ambra.demontis@unica.it>
.. moduleauthor:: Marco Melis <marco.melis@unica.it>

"""
import pickle
import gzip

from secml.utils import fm

# Remember to add any new method to following list
__all__ = ['save', 'load']


[docs]def save(file_path, obj): """Save object to file using cPickle. This functions stores a generic python object into a compressed gzip file (`*.gz`). Saved objects can be loaded using `.load`. Parameters ---------- file_path : str Path to destination file. obj : object Any python object to save. Returns ------- obj_path : str Full path to the stored object. Notes ----- Objects are stored using **protocol 4** data stream format. For more information see https://docs.python.org/3/library/pickle.html#data-stream-format """ # Adding extension to destination file if user forgot about it... file_ext = fm.splitext(file_path)[1] file_path = file_path + '.gz' if file_ext != '.gz' else file_path # open the reference to target file with gzip.open(file_path, 'wb') as f_ref: pickle.dump(obj, f_ref, protocol=4) return fm.join(fm.abspath(file_path), fm.split(file_path)[1])
[docs]def load(file_path, encoding='bytes'): """Load object from cPickle file. Load a generic gzip compressed python object created by `.save`. Parameters ---------- file_path : str Path to target file to read. encoding : str, optional Encoding to use for loading the file. Default 'bytes'. """ with gzip.open(file_path, 'rb') as f_ref: # Loading and returning the object return pickle.load(f_ref, encoding=encoding)