From 31af3084ab941b96ca9536e40aaaab5b0194ed54 Mon Sep 17 00:00:00 2001 From: Matthew Roeschke Date: Fri, 1 Feb 2019 10:50:35 -0800 Subject: [PATCH] CLN: to_pickle internals (#25044) --- pandas/compat/pickle_compat.py | 3 +- pandas/io/pickle.py | 73 ++++++++++------------------------ 2 files changed, 22 insertions(+), 54 deletions(-) diff --git a/pandas/compat/pickle_compat.py b/pandas/compat/pickle_compat.py index 61295b8249f58b..8f16f8154b952e 100644 --- a/pandas/compat/pickle_compat.py +++ b/pandas/compat/pickle_compat.py @@ -201,7 +201,7 @@ def load_newobj_ex(self): pass -def load(fh, encoding=None, compat=False, is_verbose=False): +def load(fh, encoding=None, is_verbose=False): """load a pickle, with a provided encoding if compat is True: @@ -212,7 +212,6 @@ def load(fh, encoding=None, compat=False, is_verbose=False): ---------- fh : a filelike object encoding : an optional encoding - compat : provide Series compatibility mode, boolean, default False is_verbose : show exception output """ diff --git a/pandas/io/pickle.py b/pandas/io/pickle.py index 789f55a62dc58c..ab4a266853a78a 100644 --- a/pandas/io/pickle.py +++ b/pandas/io/pickle.py @@ -1,8 +1,7 @@ """ pickle compat """ import warnings -import numpy as np -from numpy.lib.format import read_array, write_array +from numpy.lib.format import read_array from pandas.compat import PY3, BytesIO, cPickle as pkl, pickle_compat as pc @@ -76,6 +75,7 @@ def to_pickle(obj, path, compression='infer', protocol=pkl.HIGHEST_PROTOCOL): try: f.write(pkl.dumps(obj, protocol=protocol)) finally: + f.close() for _f in fh: _f.close() @@ -138,63 +138,32 @@ def read_pickle(path, compression='infer'): >>> os.remove("./dummy.pkl") """ path = _stringify_path(path) + f, fh = _get_handle(path, 'rb', compression=compression, is_text=False) + + # 1) try with cPickle + # 2) try with the compat pickle to handle subclass changes + # 3) pass encoding only if its not None as py2 doesn't handle the param - def read_wrapper(func): - # wrapper file handle open/close operation - f, fh = _get_handle(path, 'rb', - compression=compression, - is_text=False) - try: - return func(f) - finally: - for _f in fh: - _f.close() - - def try_read(path, encoding=None): - # try with cPickle - # try with current pickle, if we have a Type Error then - # try with the compat pickle to handle subclass changes - # pass encoding only if its not None as py2 doesn't handle - # the param - - # cpickle - # GH 6899 - try: - with warnings.catch_warnings(record=True): - # We want to silence any warnings about, e.g. moved modules. - warnings.simplefilter("ignore", Warning) - return read_wrapper(lambda f: pkl.load(f)) - except Exception: # noqa: E722 - # reg/patched pickle - # compat not used in pandas/compat/pickle_compat.py::load - # TODO: remove except block OR modify pc.load to use compat - try: - return read_wrapper( - lambda f: pc.load(f, encoding=encoding, compat=False)) - # compat pickle - except Exception: # noqa: E722 - return read_wrapper( - lambda f: pc.load(f, encoding=encoding, compat=True)) try: - return try_read(path) + with warnings.catch_warnings(record=True): + # We want to silence any warnings about, e.g. moved modules. + warnings.simplefilter("ignore", Warning) + return pkl.load(f) except Exception: # noqa: E722 - if PY3: - return try_read(path, encoding='latin1') - raise - + try: + return pc.load(f, encoding=None) + except Exception: # noqa: E722 + if PY3: + return pc.load(f, encoding='latin1') + raise + finally: + f.close() + for _f in fh: + _f.close() # compat with sparse pickle / unpickle -def _pickle_array(arr): - arr = arr.view(np.ndarray) - - buf = BytesIO() - write_array(buf, arr) - - return buf.getvalue() - - def _unpickle_array(bytes): arr = read_array(BytesIO(bytes))