MemoryRecords iterator; MemoryRecordsBuilder records() helper

Author: dpkp

kafka/record/memory_records.py

@@ -109,6 +109,16 @@ def next_batch(self, _min_slice=MIN_SLICE,
         else:
             return DefaultRecordBatch(next_slice)
 
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if not self.has_next():
+            raise StopIteration
+        return self.next_batch()
+
+    next = __next__
+
 
 class MemoryRecordsBuilder(object):
 
@@ -186,6 +196,10 @@ def producer_id(self):
     def producer_epoch(self):
         return self._producer_epoch
 
+    def records(self):
+        assert self._closed
+        return MemoryRecords(self._buffer)
+
     def close(self):
         # This method may be called multiple times on the same batch
         # i.e., on retries