add compression support for pickle

Author: goldenbull

pandas/core/generic.py

@@ -1256,17 +1256,19 @@ def to_sql(self, name, con, flavor=None, schema=None, if_exists='fail',
                    if_exists=if_exists, index=index, index_label=index_label,
                    chunksize=chunksize, dtype=dtype)
 
-    def to_pickle(self, path):
+    def to_pickle(self, path, compression='infer'):
         """
         Pickle (serialize) object to input file path.
 
         Parameters
         ----------
         path : string
             File path
+        compression : {'infer', 'gzip', 'bz2', 'xz', None}, default 'infer'
+            .. versionadded:: 0.19.2
         """
         from pandas.io.pickle import to_pickle
-        return to_pickle(self, path)
+        return to_pickle(self, path, compression=compression)
 
     def to_clipboard(self, excel=None, sep=None, **kwargs):
         """

pandas/io/common.py

@@ -296,7 +296,7 @@ def _infer_compression(filepath_or_buffer, compression):
 
 
 def _get_handle(path_or_buf, mode, encoding=None, compression=None,
-                memory_map=False):
+                memory_map=False, is_text=True):
     """
     Get file handle for given path/buffer and mode.
 
@@ -311,7 +311,8 @@ def _get_handle(path_or_buf, mode, encoding=None, compression=None,
         Supported compression protocols are gzip, bz2, zip, and xz
     memory_map : boolean, default False
         See parsers._parser_params for more information.
-
+    is_text : boolean, default True
+        whether file/buffer is in text format (csv, json, etc.), or in binary mode (pickle, etc.)
     Returns
     -------
     f : file-like
@@ -385,13 +386,16 @@ def _get_handle(path_or_buf, mode, encoding=None, compression=None,
         elif encoding:
             # Python 3 and encoding
             f = open(path_or_buf, mode, encoding=encoding)
-        else:
+        elif is_text:
             # Python 3 and no explicit encoding
             f = open(path_or_buf, mode, errors='replace')
+        else:
+            # Python 3 and binary mode
+            f = open(path_or_buf, mode)
         handles.append(f)
 
     # in Python 3, convert BytesIO or fileobjects passed with an encoding
-    if compat.PY3 and (compression or isinstance(f, compat.BytesIO)):
+    if compat.PY3 and is_text and (compression or isinstance(f, compat.BytesIO)):
         from io import TextIOWrapper
         f = TextIOWrapper(f, encoding=encoding)
         handles.append(f)

pandas/io/pickle.py

@@ -4,9 +4,10 @@
 from numpy.lib.format import read_array, write_array
 from pandas.compat import BytesIO, cPickle as pkl, pickle_compat as pc, PY3
 from pandas.types.common import is_datetime64_dtype, _NS_DTYPE
+from pandas.io.common import _get_handle, _infer_compression
 
 
-def to_pickle(obj, path):
+def to_pickle(obj, path, compression='infer'):
     """
     Pickle (serialize) object to input file path
 
@@ -15,12 +16,16 @@ def to_pickle(obj, path):
     obj : any object
     path : string
         File path
+    compression : {'infer', 'gzip', 'bz2', 'xz', None}, default 'infer'
+        .. versionadded:: 0.19.2
     """
-    with open(path, 'wb') as f:
+    inferred_compression = _infer_compression(path, compression)
+    f, fh = _get_handle(path, 'wb', compression=inferred_compression, is_text=False)
+    with f:
         pkl.dump(obj, f, protocol=pkl.HIGHEST_PROTOCOL)
 
 
-def read_pickle(path):
+def read_pickle(path, compression='infer'):
     """
     Load pickled pandas object (or any other pickled object) from the specified
     file path
@@ -32,12 +37,16 @@ def read_pickle(path):
     ----------
     path : string
         File path
+    compression : {'infer', 'gzip', 'bz2', 'xz', None}, default 'infer'
+        .. versionadded:: 0.19.2
 
     Returns
     -------
     unpickled : type of object stored in file
     """
 
+    inferred_compression = _infer_compression(path, compression)
+
     def try_read(path, encoding=None):
         # try with cPickle
         # try with current pickle, if we have a Type Error then
@@ -48,18 +57,21 @@ def try_read(path, encoding=None):
         # cpickle
         # GH 6899
         try:
-            with open(path, 'rb') as fh:
-                return pkl.load(fh)
+            f, fh = _get_handle(path, 'rb', compression=inferred_compression, is_text=False)
+            with f:
+                return pkl.load(f)
         except Exception:
             # reg/patched pickle
             try:
-                with open(path, 'rb') as fh:
-                    return pc.load(fh, encoding=encoding, compat=False)
+                f, fh = _get_handle(path, 'rb', compression=inferred_compression, is_text=False)
+                with f:
+                    return pc.load(f, encoding=encoding, compat=False)
 
             # compat pickle
             except:
-                with open(path, 'rb') as fh:
-                    return pc.load(fh, encoding=encoding, compat=True)
+                f, fh = _get_handle(path, 'rb', compression=inferred_compression, is_text=False)
+                with f:
+                    return pc.load(f, encoding=encoding, compat=True)
 
     try:
         return try_read(path)
@@ -68,6 +80,7 @@ def try_read(path, encoding=None):
             return try_read(path, encoding='latin1')
         raise
 
+
 # compat with sparse pickle / unpickle
 
 

pandas/io/tests/test_pickle.py

@@ -284,6 +284,40 @@ def test_pickle_v0_15_2(self):
         #
         tm.assert_categorical_equal(cat, pd.read_pickle(pickle_path))
 
+    def compression_explicit(self, compression):
+        # issue 11666
+        with tm.ensure_clean(self.path) as path:
+            df = tm.makeDataFrame()
+            df.to_pickle(path, compression=compression)
+            tm.assert_frame_equal(df, pandas.read_pickle(path, compression=compression))
+
+    def test_compression_explicit(self):
+        compressions = [None, 'gzip', 'bz2', 'xz']
+        for c in compressions:
+            yield self.compression_explicit, c
+
+    def compression_explicit_bad(self, compression):
+        with tm.assertRaisesRegexp(ValueError, "Unrecognized compression type"):
+            with tm.ensure_clean(self.path) as path:
+                df = tm.makeDataFrame()
+                df.to_pickle(path, compression=compression)
+
+    def test_compression_explicit_bad(self):
+        compressions = ['', 'None', 'bad', '7z']
+        for c in compressions:
+            yield self.compression_explicit_bad, c
+
+    def compression_infer(self, ext):
+        with tm.ensure_clean(self.path + ext) as p:
+            df = tm.makeDataFrame()
+            df.to_pickle(p)
+            tm.assert_frame_equal(df, pandas.read_pickle(p))
+
+    def test_compression_infer(self):
+        extensions = ['', '.gz', '.bz2', '.xz', '.who_am_i']
+        for ext in extensions:
+            yield self.compression_infer, ext
+
 
 if __name__ == '__main__':
     nose.runmodule(argv=[__file__, '-vvs', '-x', '--pdb', '--pdb-failure'],