Added tests for new 'protocol' parameter in 'to_pickle'.

Author: jbschiratti

pandas/tests/io/test_pickle.py

@@ -20,11 +20,12 @@
 from distutils.version import LooseVersion
 import pandas as pd
 from pandas import Index
-from pandas.compat import is_platform_little_endian
+from pandas.compat import is_platform_little_endian, cPickle as pkl
 import pandas
 import pandas.util.testing as tm
 from pandas.tseries.offsets import Day, MonthEnd
 import shutil
+import sys
 
 
 @pytest.fixture(scope='module')
@@ -489,3 +490,37 @@ def test_read_infer(self, ext, get_random_path):
             df2 = pd.read_pickle(p2)
 
             tm.assert_frame_equal(df, df2)
+
+
+# ---------------------
+# test pickle compression
+# ---------------------
+
+class TestProtocol(object):
+
+    @pytest.mark.parametrize('protocol', [-1, 0, 1, 2])
+    def test_read(self, protocol, get_random_path):
+        with tm.ensure_clean(get_random_path) as path:
+            df = tm.makeDataFrame()
+            df.to_pickle(path, protocol=protocol)
+            df2 = pd.read_pickle(path)
+            tm.assert_frame_equal(df, df2)
+
+    @pytest.mark.parametrize('protocol', [3, 4])
+    @pytest.mark.skipif(sys.version_info[:2] >= (3, 4),
+                        reason="Testing invalid parameters for "
+                               "Python 2.x and 3.y (y < 4).")
+    def test_read_bad_versions(self, protocol, get_random_path):
+        # For Python 2.x (respectively 3.y with y < 4), [expected]
+        # HIGHEST_PROTOCOL should be 2 (respectively 3). Hence, the protocol
+        # parameter should not exceed 2 (respectively 3).
+        if sys.version_info[:2] < (3, 0):
+            expect_hp = 2
+        else:
+            expect_hp = 3
+        with tm.assert_raises_regex(ValueError,
+                                    "pickle protocol must be <= %d" %
+                                    expect_hp):
+            with tm.ensure_clean(get_random_path) as path:
+                df = tm.makeDataFrame()
+                df.to_pickle(path, protocol=protocol)