def RecordReader(stream, tolerate_pre_error=True, tolerate_subsequent_error=False):
got_first = False
reader = TransactionReader(stream)
while True:
# wind reader to just past first sentinel
while True:
with reader:
block = reader.read(READ_BUFFER_SIZE)
if not block:
return # eof
match = find_sentinel(block)
if not match:
DumpedData(got_first, tolerate_pre_error, tolerate_subsequent_error)
reader.commit()
else:
if match.start() > 0:
DumpedData(got_first, tolerate_pre_error, tolerate_subsequent_error)
reader.commit(match.end())
break
# do the read
with reader:
try:
length = cPickle.load(reader)
data = reader.read(length)
data_hash = cPickle.load(reader)
if data_hash != hash(data) or length != len(data):
raise _HashFail
data = fix_read(data)
reader.commit()
got_first = True
yield data
except (EOFError, cPickle.UnpicklingError, _HashFail):
DumpedData(got_first, tolerate_pre_error, tolerate_subsequent_error)
def test_transaction_reader(self):
sio = StringIO()
sio.write('1234567890abcdefg')
sio.seek(0)
reader = TransactionReader(sio)
# no commit, unwind read
with reader:
self.assertEqual(reader.read(2), '12')
self.assertEqual(reader.read(5), '34567')
self.assertEqual(reader.mem_use(), 7)
# commit
with reader:
self.assertEqual(reader.read(2), '12')
self.assertEqual(reader.mem_use(), 7)
self.assertEqual(reader.read(5), '34567')
self.assertEqual(reader.read(2), '89')
reader.commit()
self.assertEqual(reader.mem_use(), 0)
# partial commit
with reader:
self.assertEqual(reader.read(3), '0ab')
self.assertEqual(reader.mem_use(), 3)
reader.commit(2)
self.assertEqual(reader.mem_use(), 1) # just 'b' left
self.assertEqual(reader.read(2), 'cd')
with reader:
self.assertEqual(reader.read(2), 'bc') # '0a' was committed
self.assertEqual(reader.mem_use(), 3)
with reader:
self.assertEqual(reader.read(2), 'bc')
self.assertEqual(reader.mem_use(), 3)
self.assertEqual(reader.mem_use(), 3)
# read to end
with reader:
self.assertEqual(reader.read(), 'bcdefg')
self.assertEqual(reader.mem_use(), 6)
with reader:
self.assertEqual(reader.read(), 'bcdefg')
reader.commit()
self.assertEqual(reader.mem_use(), 0)
# eof
with reader:
self.assertEqual(reader.read(3), '') # EOF
self.assertEqual(reader.read(), '')
# newline and read-line
sio = StringIO()
sio.write('abc123\nanotherline\nlastlinewith no eof')
sio.seek(0)
reader = TransactionReader(sio)
with reader:
self.assertEqual(reader.readline(), 'abc123\n')
reader.commit()
with reader:
self.assertEqual(reader.readline(), 'anotherline\n')
with reader:
self.assertEqual(reader.readline(), 'anotherline\n')
self.assertEqual(reader.readline(), 'lastlinewith no eof')
reader.commit()
with reader:
self.assertEqual(reader.readline(), '')
# read a lot, roll-back, read a little, commit all, verify
sio.seek(0)
reader = TransactionReader(sio)
with reader:
reader.read(10)
self.assertEqual(reader.mem_use(), 10)
with reader:
reader.read(2)
reader.commit()
self.assertEqual(reader.mem_use(), 8)
