Skip to content

Commit c4637c8

Browse files
mroeschkePingviinituutti
authored andcommitted
CLN: to_pickle internals (pandas-dev#25044)
1 parent ea0410a commit c4637c8

File tree

2 files changed

+22
-54
lines changed

2 files changed

+22
-54
lines changed

pandas/compat/pickle_compat.py

+1-2
Original file line numberDiff line numberDiff line change
@@ -201,7 +201,7 @@ def load_newobj_ex(self):
201201
pass
202202

203203

204-
def load(fh, encoding=None, compat=False, is_verbose=False):
204+
def load(fh, encoding=None, is_verbose=False):
205205
"""load a pickle, with a provided encoding
206206
207207
if compat is True:
@@ -212,7 +212,6 @@ def load(fh, encoding=None, compat=False, is_verbose=False):
212212
----------
213213
fh : a filelike object
214214
encoding : an optional encoding
215-
compat : provide Series compatibility mode, boolean, default False
216215
is_verbose : show exception output
217216
"""
218217

pandas/io/pickle.py

+21-52
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,7 @@
11
""" pickle compat """
22
import warnings
33

4-
import numpy as np
5-
from numpy.lib.format import read_array, write_array
4+
from numpy.lib.format import read_array
65

76
from pandas.compat import PY3, BytesIO, cPickle as pkl, pickle_compat as pc
87

@@ -76,6 +75,7 @@ def to_pickle(obj, path, compression='infer', protocol=pkl.HIGHEST_PROTOCOL):
7675
try:
7776
f.write(pkl.dumps(obj, protocol=protocol))
7877
finally:
78+
f.close()
7979
for _f in fh:
8080
_f.close()
8181

@@ -138,63 +138,32 @@ def read_pickle(path, compression='infer'):
138138
>>> os.remove("./dummy.pkl")
139139
"""
140140
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
141146

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))
178147
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)
180152
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()
185163

186164
# compat with sparse pickle / unpickle
187165

188166

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-
198167
def _unpickle_array(bytes):
199168
arr = read_array(BytesIO(bytes))
200169

0 commit comments

Comments
 (0)