FIX: Fix DBNStore.from_bytes with seekable buffer

Author: nmacholl

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## 0.21.0 - TBD
+
+#### Bug fixes
+- Fixed an issue where `DBNStore.from_bytes` did not rewind seekable buffers
+
 ## 0.20.0 - 2023-09-21
 
 #### Enhancements

databento/common/dbnstore.py

@@ -252,7 +252,12 @@ class MemoryDataSource(DataSource):
     """
 
     def __init__(self, source: BytesIO | bytes | IO[bytes]):
-        initial_data = source if isinstance(source, bytes) else source.read()
+        if isinstance(source, bytes):
+            initial_data = source
+        else:
+            source.seek(0)
+            initial_data = source.read()
+
         if len(initial_data) == 0:
             raise ValueError(
                 f"Cannot create data source from empty {type(source).__name__}",

tests/test_historical_bento.py

@@ -924,6 +924,24 @@ def test_dbnstore_buffer_long(
     with pytest.raises(BentoError):
         dbnstore.to_json(tmp_path / "test.json")
 
+def test_dbnstore_buffer_rewind(
+    test_data: Callable[[Schema], bytes],
+    tmp_path: Path,
+) -> None:
+    """
+    Test that creating a DBNStore from a seekable buffer will rewind.
+    """
+    # Arrange
+    dbn_stub_data = (
+        zstandard.ZstdDecompressor().stream_reader(test_data(Schema.MBO)).read()
+    )
+
+    # Act
+    dbn_bytes = BytesIO()
+    dbn_bytes.write(dbn_stub_data)
+    dbnstore = DBNStore.from_bytes(data=dbn_bytes)
+
+    assert len(dbnstore.to_df()) == 4
 
 @pytest.mark.parametrize(
     "schema",