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)