|
1 | 1 | """ pickle compat """
|
2 | 2 | import warnings
|
3 | 3 |
|
4 |
| -import numpy as np |
5 |
| -from numpy.lib.format import read_array, write_array |
| 4 | +from numpy.lib.format import read_array |
6 | 5 |
|
7 | 6 | from pandas.compat import PY3, BytesIO, cPickle as pkl, pickle_compat as pc
|
8 | 7 |
|
@@ -76,6 +75,7 @@ def to_pickle(obj, path, compression='infer', protocol=pkl.HIGHEST_PROTOCOL):
|
76 | 75 | try:
|
77 | 76 | f.write(pkl.dumps(obj, protocol=protocol))
|
78 | 77 | finally:
|
| 78 | + f.close() |
79 | 79 | for _f in fh:
|
80 | 80 | _f.close()
|
81 | 81 |
|
@@ -138,63 +138,32 @@ def read_pickle(path, compression='infer'):
|
138 | 138 | >>> os.remove("./dummy.pkl")
|
139 | 139 | """
|
140 | 140 | path = _stringify_path(path)
|
| 141 | + f, fh = _get_handle(path, 'rb', compression=compression, is_text=False) |
| 142 | + |
| 143 | + # 1) try with cPickle |
| 144 | + # 2) try with the compat pickle to handle subclass changes |
| 145 | + # 3) pass encoding only if its not None as py2 doesn't handle the param |
141 | 146 |
|
142 |
| - def read_wrapper(func): |
143 |
| - # wrapper file handle open/close operation |
144 |
| - f, fh = _get_handle(path, 'rb', |
145 |
| - compression=compression, |
146 |
| - is_text=False) |
147 |
| - try: |
148 |
| - return func(f) |
149 |
| - finally: |
150 |
| - for _f in fh: |
151 |
| - _f.close() |
152 |
| - |
153 |
| - def try_read(path, encoding=None): |
154 |
| - # try with cPickle |
155 |
| - # try with current pickle, if we have a Type Error then |
156 |
| - # try with the compat pickle to handle subclass changes |
157 |
| - # pass encoding only if its not None as py2 doesn't handle |
158 |
| - # the param |
159 |
| - |
160 |
| - # cpickle |
161 |
| - # GH 6899 |
162 |
| - try: |
163 |
| - with warnings.catch_warnings(record=True): |
164 |
| - # We want to silence any warnings about, e.g. moved modules. |
165 |
| - warnings.simplefilter("ignore", Warning) |
166 |
| - return read_wrapper(lambda f: pkl.load(f)) |
167 |
| - except Exception: # noqa: E722 |
168 |
| - # reg/patched pickle |
169 |
| - # compat not used in pandas/compat/pickle_compat.py::load |
170 |
| - # TODO: remove except block OR modify pc.load to use compat |
171 |
| - try: |
172 |
| - return read_wrapper( |
173 |
| - lambda f: pc.load(f, encoding=encoding, compat=False)) |
174 |
| - # compat pickle |
175 |
| - except Exception: # noqa: E722 |
176 |
| - return read_wrapper( |
177 |
| - lambda f: pc.load(f, encoding=encoding, compat=True)) |
178 | 147 | try:
|
179 |
| - return try_read(path) |
| 148 | + with warnings.catch_warnings(record=True): |
| 149 | + # We want to silence any warnings about, e.g. moved modules. |
| 150 | + warnings.simplefilter("ignore", Warning) |
| 151 | + return pkl.load(f) |
180 | 152 | except Exception: # noqa: E722
|
181 |
| - if PY3: |
182 |
| - return try_read(path, encoding='latin1') |
183 |
| - raise |
184 |
| - |
| 153 | + try: |
| 154 | + return pc.load(f, encoding=None) |
| 155 | + except Exception: # noqa: E722 |
| 156 | + if PY3: |
| 157 | + return pc.load(f, encoding='latin1') |
| 158 | + raise |
| 159 | + finally: |
| 160 | + f.close() |
| 161 | + for _f in fh: |
| 162 | + _f.close() |
185 | 163 |
|
186 | 164 | # compat with sparse pickle / unpickle
|
187 | 165 |
|
188 | 166 |
|
189 |
| -def _pickle_array(arr): |
190 |
| - arr = arr.view(np.ndarray) |
191 |
| - |
192 |
| - buf = BytesIO() |
193 |
| - write_array(buf, arr) |
194 |
| - |
195 |
| - return buf.getvalue() |
196 |
| - |
197 |
| - |
198 | 167 | def _unpickle_array(bytes):
|
199 | 168 | arr = read_array(BytesIO(bytes))
|
200 | 169 |
|
|
0 commit comments