Skip to content

Commit

Permalink
Fix-up test_pickle_empty (#5303)
Browse files Browse the repository at this point in the history
Switches to using the internal utility class `MemoryviewHolder` instead of NumPy to allow testing without NumPy. Calls `serialize` directly with `pickle` as the only option (instead of calling `pickle_dumps`). Though also `assert`s pickling did occur. Also includes some other minor fixes.

Only tries manipulating writability of frames with pickle protocol 5 in use. This doesn't work for earlier pickle protocols as there is only one `bytes` object containing everything (and no way to inspect what went into it, let alone reproduce it).
  • Loading branch information
jakirkham committed Sep 10, 2021
1 parent 31afb54 commit 0774365
Showing 1 changed file with 33 additions and 17 deletions.
50 changes: 33 additions & 17 deletions distributed/protocol/tests/test_pickle.py
Expand Up @@ -8,7 +8,6 @@

from distributed.protocol import deserialize, serialize
from distributed.protocol.pickle import HIGHEST_PROTOCOL, dumps, loads
from distributed.protocol.serialize import pickle_dumps

if sys.version_info < (3, 8):
try:
Expand All @@ -19,6 +18,17 @@
import pickle


class MemoryviewHolder:
def __init__(self, mv):
self.mv = memoryview(mv)

def __reduce_ex__(self, protocol):
if protocol >= 5:
return MemoryviewHolder, (pickle.PickleBuffer(self.mv),)
else:
return MemoryviewHolder, (self.mv.tobytes(),)


def test_pickle_data():
data = [1, b"123", "123", [123], {}, set()]
for d in data:
Expand All @@ -27,16 +37,6 @@ def test_pickle_data():


def test_pickle_out_of_band():
class MemoryviewHolder:
def __init__(self, mv):
self.mv = memoryview(mv)

def __reduce_ex__(self, protocol):
if protocol >= 5:
return MemoryviewHolder, (pickle.PickleBuffer(self.mv),)
else:
return MemoryviewHolder, (self.mv.tobytes(),)

mv = memoryview(b"123")
mvh = MemoryviewHolder(mv)

Expand Down Expand Up @@ -73,13 +73,29 @@ def __reduce_ex__(self, protocol):


def test_pickle_empty():
np = pytest.importorskip("numpy")
x = np.arange(2)[0:0] # Empty view
header, frames = pickle_dumps(x)
header["writeable"] = [False] * len(frames)
x = MemoryviewHolder(bytearray()) # Empty view
header, frames = serialize(x, serializers=("pickle",))

assert header["serializer"] == "pickle"
assert len(frames) >= 1
assert isinstance(frames[0], bytes)

if HIGHEST_PROTOCOL >= 5:
assert len(frames) == 2
assert len(header["writeable"]) == 1

header["writeable"] = (False,) * len(frames)
else:
assert len(frames) == 1
assert len(header["writeable"]) == 0

y = deserialize(header, frames)
assert memoryview(y).nbytes == 0
assert memoryview(y).readonly

assert isinstance(y, MemoryviewHolder)
assert isinstance(y.mv, memoryview)
assert y.mv == x.mv
assert y.mv.nbytes == 0
assert y.mv.readonly


def test_pickle_numpy():
Expand Down

0 comments on commit 0774365

Please sign in to comment.