MOD: Implement Iterable for Bento class

Author: nmacholl

databento/common/bento.py

@@ -2,10 +2,21 @@
 
 import abc
 import datetime as dt
+import logging
 from io import BytesIO
 from os import PathLike
 from pathlib import Path
-from typing import IO, TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
+from typing import (
+    IO,
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Generator,
+    List,
+    Optional,
+    Union,
+)
 
 import numpy as np
 import pandas as pd
@@ -23,6 +34,8 @@
 from databento.common.symbology import ProductIdMappingInterval
 
 
+logger = logging.getLogger(__name__)
+
 if TYPE_CHECKING:
     from databento.historical.client import Historical
 
@@ -274,6 +287,14 @@ def __init__(self, data_source: DataSource) -> None:
             Dict[int, str],
         ] = {}
 
+    def __iter__(self) -> Generator[np.void, None, None]:
+        for _ in range(self.record_count):
+            raw = self.reader.read(self.record_size)
+            rec = np.frombuffer(raw, dtype=STRUCT_MAP[self.schema])
+            if rec.size == 0:
+                raise StopIteration
+            yield rec[0]
+
     def _apply_pretty_ts(self, df: pd.DataFrame) -> pd.DataFrame:
         df.index = pd.to_datetime(df.index, utc=True)
         for column in df.columns:
@@ -672,14 +693,15 @@ def replay(self, callback: Callable[[Any], None]) -> None:
             The callback to the data handler.
 
         """
-        dtype = STRUCT_MAP[self.schema]
-        reader: IO[bytes] = self.reader
-        while True:
-            raw: bytes = reader.read(self.record_size)
-            record = np.frombuffer(raw, dtype=dtype)
-            if record.size == 0:
-                break
-            callback(record[0])
+        for record in self:
+            try:
+                callback(record)
+            except Exception as exc:
+                logger.exception(
+                    "exception while replaying to user callback",
+                    exc_info=exc,
+                )
+                raise
 
     def request_full_definitions(
         self,

tests/test_historical_bento.py

@@ -598,3 +598,62 @@ def test_mbp_1_to_json_with_all_options_writes_expected_file_to_disk(self) -> No
 
         # Cleanup
         os.remove(path)
+
+    def test_bento_iterable(self) -> None:
+        """
+        Tests the Bento iterable implementation to ensure records
+        can be accessed by iteration.
+        """
+        # Arrange
+        stub_data = get_test_data(schema=Schema.MBO)
+        bento = Bento.from_bytes(data=stub_data)
+
+        record_list = list(bento)
+        assert (
+            str(record_list[0])
+            == "(14, 160, 1, 5482, 1609160400000429831, 647784973705, "
+            "3722750000000, 1, -128, 0, b'C', b'A', 1609160400000704060, "
+            "22993, 1170352)"
+        )
+        assert (
+            str(record_list[1])
+            == "(14, 160, 1, 5482, 1609160400000429831, 647784973705, "
+            "3722750000000, 1, -128, 0, b'C', b'A', 1609160400000704060, "
+            "22993, 1170352)"
+        )
+
+    def test_bento_iterable_parallel(self) -> None:
+        """
+        Tests the Bento iterable implementation to ensure iterators are
+        not stateful. For example, calling next() on one iterator does
+        not affect another.
+        """
+        # Arrange
+        stub_data = get_test_data(schema=Schema.MBO)
+        bento = Bento.from_bytes(data=stub_data)
+
+        first = iter(bento)
+        second = iter(bento)
+
+        assert (
+            str(next(first)) == "(14, 160, 1, 5482, 1609160400000429831, 647784973705, "
+            "3722750000000, 1, -128, 0, b'C', b'A', 1609160400000704060, "
+            "22993, 1170352)"
+        )
+        assert (
+            str(next(second))
+            == "(14, 160, 1, 5482, 1609160400000429831, 647784973705, "
+            "3722750000000, 1, -128, 0, b'C', b'A', 1609160400000704060, "
+            "22993, 1170352)"
+        )
+        assert (
+            str(next(second))
+            == "(14, 160, 1, 5482, 1609160400000429831, 647784973705, "
+            "3722750000000, 1, -128, 0, b'C', b'A', 1609160400000704060, "
+            "22993, 1170352)"
+        )
+        assert (
+            str(next(first)) == "(14, 160, 1, 5482, 1609160400000429831, 647784973705, "
+            "3722750000000, 1, -128, 0, b'C', b'A', 1609160400000704060, "
+            "22993, 1170352)"
+        )