def __init__(self, connection):
self.sio = BytesIO()
self.pickler = Pickler(self.sio, PICKLE_PROTOCOL)
self.pickler.persistent_id = method(call_if_persistent,
self._persistent_id)
self.objects_found = []
self.refs = set() # populated by _persistent_id()
self.connection = connection
def get_state(self, data, load=True):
self.load_count += 1
s = BytesIO()
s.write(data)
s.seek(0)
unpickler = self._get_unpickler(s)
klass = unpickler.load()
position = s.tell()
if data[s.tell()] == COMPRESSED_START_BYTE:
# This is almost certainly a compressed pickle.
try:
decompressed = decompress(data[position:])
except zlib_error:
pass # let the unpickler try anyway.
else:
s.write(decompressed)
s.seek(position)
if load:
return unpickler.load()
else:
return s.read()
def d(self):
s = BytesIO(as_bytes('nope'))
assert raises(AssertionError, Shelf, s, readonly=True)
s = BytesIO(as_bytes("SHELF-1\nbogus"))
assert raises(ShortRead, Shelf, s, readonly=True)
