Skip to content

Commit

Permalink
CLN: to_pickle internals (pandas-dev#25044)
Browse files Browse the repository at this point in the history
  • Loading branch information
mroeschke authored and Pingviinituutti committed Feb 28, 2019
1 parent 06e6e60 commit 31af308
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 54 deletions.
3 changes: 1 addition & 2 deletions pandas/compat/pickle_compat.py
Original file line number Diff line number Diff line change
Expand Up @@ -201,7 +201,7 @@ def load_newobj_ex(self):
pass


def load(fh, encoding=None, compat=False, is_verbose=False):
def load(fh, encoding=None, is_verbose=False):
"""load a pickle, with a provided encoding
if compat is True:
Expand All @@ -212,7 +212,6 @@ def load(fh, encoding=None, compat=False, is_verbose=False):
----------
fh : a filelike object
encoding : an optional encoding
compat : provide Series compatibility mode, boolean, default False
is_verbose : show exception output
"""

Expand Down
73 changes: 21 additions & 52 deletions pandas/io/pickle.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,7 @@
""" pickle compat """
import warnings

import numpy as np
from numpy.lib.format import read_array, write_array
from numpy.lib.format import read_array

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

Expand Down Expand Up @@ -76,6 +75,7 @@ def to_pickle(obj, path, compression='infer', protocol=pkl.HIGHEST_PROTOCOL):
try:
f.write(pkl.dumps(obj, protocol=protocol))
finally:
f.close()
for _f in fh:
_f.close()

Expand Down Expand Up @@ -138,63 +138,32 @@ def read_pickle(path, compression='infer'):
>>> os.remove("./dummy.pkl")
"""
path = _stringify_path(path)
f, fh = _get_handle(path, 'rb', compression=compression, is_text=False)

# 1) try with cPickle
# 2) try with the compat pickle to handle subclass changes
# 3) pass encoding only if its not None as py2 doesn't handle the param

def read_wrapper(func):
# wrapper file handle open/close operation
f, fh = _get_handle(path, 'rb',
compression=compression,
is_text=False)
try:
return func(f)
finally:
for _f in fh:
_f.close()

def try_read(path, encoding=None):
# try with cPickle
# try with current pickle, if we have a Type Error then
# try with the compat pickle to handle subclass changes
# pass encoding only if its not None as py2 doesn't handle
# the param

# cpickle
# GH 6899
try:
with warnings.catch_warnings(record=True):
# We want to silence any warnings about, e.g. moved modules.
warnings.simplefilter("ignore", Warning)
return read_wrapper(lambda f: pkl.load(f))
except Exception: # noqa: E722
# reg/patched pickle
# compat not used in pandas/compat/pickle_compat.py::load
# TODO: remove except block OR modify pc.load to use compat
try:
return read_wrapper(
lambda f: pc.load(f, encoding=encoding, compat=False))
# compat pickle
except Exception: # noqa: E722
return read_wrapper(
lambda f: pc.load(f, encoding=encoding, compat=True))
try:
return try_read(path)
with warnings.catch_warnings(record=True):
# We want to silence any warnings about, e.g. moved modules.
warnings.simplefilter("ignore", Warning)
return pkl.load(f)
except Exception: # noqa: E722
if PY3:
return try_read(path, encoding='latin1')
raise

try:
return pc.load(f, encoding=None)
except Exception: # noqa: E722
if PY3:
return pc.load(f, encoding='latin1')
raise
finally:
f.close()
for _f in fh:
_f.close()

# compat with sparse pickle / unpickle


def _pickle_array(arr):
arr = arr.view(np.ndarray)

buf = BytesIO()
write_array(buf, arr)

return buf.getvalue()


def _unpickle_array(bytes):
arr = read_array(BytesIO(bytes))

Expand Down

0 comments on commit 31af308

Please sign in to comment.