|
12 | 12 | """
|
13 | 13 | import bz2
|
14 | 14 | import datetime
|
| 15 | +import functools |
15 | 16 | import glob
|
16 | 17 | import gzip
|
17 | 18 | import io
|
|
24 | 25 |
|
25 | 26 | import pytest
|
26 | 27 |
|
27 |
| -from pandas.compat import get_lzma_file, import_lzma, is_platform_little_endian |
| 28 | +from pandas.compat import PY38, get_lzma_file, import_lzma, is_platform_little_endian |
28 | 29 | import pandas.util._test_decorators as td
|
29 | 30 |
|
30 | 31 | import pandas as pd
|
@@ -155,28 +156,43 @@ def test_pickles(current_pickle_data, legacy_pickle):
|
155 | 156 | compare(current_pickle_data, legacy_pickle, version)
|
156 | 157 |
|
157 | 158 |
|
158 |
| -def test_round_trip_current(current_pickle_data): |
159 |
| - def python_pickler(obj, path): |
160 |
| - with open(path, "wb") as fh: |
161 |
| - pickle.dump(obj, fh, protocol=-1) |
| 159 | +def python_pickler(obj, path): |
| 160 | + with open(path, "wb") as fh: |
| 161 | + pickle.dump(obj, fh, protocol=-1) |
162 | 162 |
|
163 |
| - def python_unpickler(path): |
164 |
| - with open(path, "rb") as fh: |
165 |
| - fh.seek(0) |
166 |
| - return pickle.load(fh) |
167 | 163 |
|
| 164 | +def python_unpickler(path): |
| 165 | + with open(path, "rb") as fh: |
| 166 | + fh.seek(0) |
| 167 | + return pickle.load(fh) |
| 168 | + |
| 169 | + |
| 170 | +@pytest.mark.parametrize( |
| 171 | + "pickle_writer", |
| 172 | + [ |
| 173 | + pytest.param(python_pickler, id="python"), |
| 174 | + pytest.param(pd.to_pickle, id="pandas_proto_default"), |
| 175 | + pytest.param( |
| 176 | + functools.partial(pd.to_pickle, protocol=pickle.HIGHEST_PROTOCOL), |
| 177 | + id="pandas_proto_highest", |
| 178 | + ), |
| 179 | + pytest.param(functools.partial(pd.to_pickle, protocol=4), id="pandas_proto_4"), |
| 180 | + pytest.param( |
| 181 | + functools.partial(pd.to_pickle, protocol=5), |
| 182 | + id="pandas_proto_5", |
| 183 | + marks=pytest.mark.skipif(not PY38, reason="protocol 5 not supported"), |
| 184 | + ), |
| 185 | + ], |
| 186 | +) |
| 187 | +def test_round_trip_current(current_pickle_data, pickle_writer): |
168 | 188 | data = current_pickle_data
|
169 | 189 | for typ, dv in data.items():
|
170 | 190 | for dt, expected in dv.items():
|
171 | 191 |
|
172 | 192 | for writer in [pd.to_pickle, python_pickler]:
|
173 |
| - if writer is None: |
174 |
| - continue |
175 |
| - |
176 | 193 | with tm.ensure_clean() as path:
|
177 |
| - |
178 | 194 | # test writing with each pickler
|
179 |
| - writer(expected, path) |
| 195 | + pickle_writer(expected, path) |
180 | 196 |
|
181 | 197 | # test reading with each unpickler
|
182 | 198 | result = pd.read_pickle(path)
|
|
0 commit comments