CLN: to_pickle internals (#25044)

Author: mroeschke

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
     """
 

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))