MOD: Integrate databento-dbn with python client

Author: nmacholl

README.md

@@ -31,7 +31,7 @@ The library is fully compatible with the latest distribution of Anaconda 3.7 and
 The minimum dependencies as found in the `requirements.txt` are also listed below:
 - Python (>=3.7)
 - aiohttp (>=3.7.2)
-- dbz-python (>=0.2.1)
+- databento-dbn (>=0.3.2)
 - numpy (>=1.17.0)
 - pandas (>=1.1.3)
 - requests (>=2.24.0)

databento/common/bento.py

@@ -40,6 +40,48 @@
     from databento.historical.client import Historical
 
 
+def is_zstandard(reader: IO[bytes]) -> bool:
+    """
+    Determine if an `IO[bytes]` reader contains zstandard compressed
+    data.
+
+    Parameters
+    ----------
+    reader : IO[bytes]
+        The data to check.
+
+    Returns
+    -------
+    bool
+
+    """
+    reader.seek(0)  # ensure we read from the beginning
+    try:
+        zstandard.get_frame_parameters(reader.read(18))
+    except zstandard.ZstdError:
+        return False
+    else:
+        return True
+
+
+def is_dbn(reader: IO[bytes]) -> bool:
+    """
+    Determine if an `IO[bytes]` reader contains dbn data.
+
+    Parameters
+    ----------
+    reader : IO[bytes]
+        The data to check.
+
+    Returns
+    -------
+    bool
+
+    """
+    reader.seek(0)  # ensure we read from the beginning
+    return reader.read(3) == b"DBN"
+
+
 class DataSource(abc.ABC):
     """Abstract base class for backing Bento classes with data."""
 
@@ -266,19 +308,35 @@ class Bento:
     def __init__(self, data_source: DataSource) -> None:
         self._data_source = data_source
 
-        # Check for zstd skippable frame
+        # Check compression
         buffer = self._data_source.reader
-        if not buffer.read(4).startswith(b"P*M\x18"):
-            raise RuntimeError(f"{self._data_source.name} is not a valid DBN format")
+
+        if is_zstandard(buffer):
+            self._compression = Compression.ZSTD
+            buffer = zstandard.ZstdDecompressor().stream_reader(data_source.reader)
+        elif is_dbn(buffer):
+            self._compression = Compression.NONE
+            buffer = data_source.reader
+        else:
+            # We don't know how to read this file
+            raise RuntimeError(
+                f"Could not determine compression format of {self._data_source.name}",
+            )
+
+        # Get metadata length
+        metadata_bytes = BytesIO(buffer.read(8))
+        metadata_bytes.seek(4)
         metadata_length = int.from_bytes(
-            buffer.read(4),
+            metadata_bytes.read(4),
             byteorder="little",
         )
+        self._metadata_length = metadata_length + 8
 
-        buffer.seek(0)  # Rewind to read the entire header
+        metadata_bytes.write(buffer.read(metadata_length))
 
-        self._metadata: Dict[str, Any] = MetadataDecoder.decode_to_json(
-            raw_metadata=buffer.read(8 + metadata_length),
+        # Read metadata
+        self._metadata: Dict[str, Any] = MetadataDecoder().decode_to_json(
+            metadata_bytes.getvalue(),
         )
 
         # This is populated when _map_symbols is called
@@ -408,13 +466,14 @@ def _map_symbols(self, df: pd.DataFrame, pretty_ts: bool) -> pd.DataFrame:
     def compression(self) -> Compression:
         """
         Return the data compression format (if any).
+        This is determined by inspecting the data.
 
         Returns
         -------
         Compression
 
         """
-        return Compression(self._metadata["compression"])
+        return self._compression
 
     @property
     def dataset(self) -> str:
@@ -513,25 +572,37 @@ def raw(self) -> bytes:
         -------
         bytes
 
+        See Also
+        --------
+        Bento.reader
+
         """
         return self._data_source.reader.read()
 
     @property
     def reader(self) -> IO[bytes]:
         """
-        Return an I/O reader for the data.
-
-        Parameters
-        ----------
-        decompress : bool
-            If data should be decompressed.
+        Return an I/O reader for the DBN records.
 
         Returns
         -------
         BinaryIO
 
+        See Also
+        --------
+        Bento.raw
+
         """
-        return zstandard.ZstdDecompressor().stream_reader(self._data_source.reader)
+        if self.compression == Compression.ZSTD:
+            reader: IO[bytes] = zstandard.ZstdDecompressor().stream_reader(
+                self._data_source.reader,
+            )
+        else:
+            reader = self._data_source.reader
+
+        # Seek past the metadata to read records
+        reader.seek(self._metadata_length)
+        return reader
 
     @property
     def record_count(self) -> int:

databento/common/enums.py

@@ -154,38 +154,6 @@ class Schema(StringyMixin, str, Enum):
     GATEWAY_ERROR = "gateway_error"
     SYMBOL_MAPPING = "symbol_mapping"
 
-    @classmethod
-    def from_int(cls, value: int) -> "Schema":
-        if value == 0:
-            return cls.MBO
-        if value == 1:
-            return cls.MBP_1
-        if value == 2:
-            return cls.MBP_10
-        if value == 3:
-            return cls.TBBO
-        if value == 4:
-            return cls.TRADES
-        if value == 5:
-            return cls.OHLCV_1S
-        if value == 6:
-            return cls.OHLCV_1M
-        if value == 7:
-            return cls.OHLCV_1H
-        if value == 8:
-            return cls.OHLCV_1D
-        if value == 9:
-            return cls.DEFINITION
-        if value == 10:
-            return cls.STATISTICS
-        if value == 11:
-            return cls.STATUS
-        if value == 12:
-            return cls.GATEWAY_ERROR
-        if value == 13:
-            return cls.SYMBOL_MAPPING
-        raise ValueError(f"value `{value}` is not a valid member of {cls.__name__}")
-
 
 @unique
 @coercible
@@ -205,14 +173,6 @@ class Compression(StringyMixin, str, Enum):
     NONE = "none"
     ZSTD = "zstd"
 
-    @classmethod
-    def from_int(cls, value: int) -> "Compression":
-        if value == 0:
-            return cls.NONE
-        if value == 1:
-            return cls.ZSTD
-        raise ValueError(f"value `{value}` is not a valid member of {cls.__name__}")
-
 
 @unique
 @coercible
@@ -254,16 +214,6 @@ class SType(StringyMixin, str, Enum):
     NATIVE = "native"
     SMART = "smart"
 
-    @classmethod
-    def from_int(cls, value: int) -> "SType":
-        if value == 0:
-            return cls.PRODUCT_ID
-        if value == 1:
-            return cls.NATIVE
-        if value == 2:
-            return cls.SMART
-        raise ValueError(f"value `{value}` is not a valid member of {cls.__name__}")
-
 
 @unique
 @coercible

databento/common/metadata.py

@@ -1,7 +1,6 @@
 from typing import Dict
 
-from databento.common.enums import Compression, Schema, SType
-from dbz_python import decode_metadata
+from databento_dbn import decode_metadata
 
 
 class MetadataDecoder:
@@ -26,15 +25,17 @@ def decode_to_json(raw_metadata: bytes) -> Dict[str, object]:
         """
         metadata = decode_metadata(raw_metadata)
         conversion_mapping = {
-            "compression": lambda c: str(Compression.from_int(c)),
             "limit": lambda lim: None if lim == 0 else lim,
-            "mappings": lambda m: {i["native"]: i["intervals"] for i in m},
-            "schema": lambda s: str(Schema.from_int(s)),
-            "stype_in": lambda s_in: str(SType.from_int(s_in)),
-            "stype_out": lambda s_out: str(SType.from_int(s_out)),
         }
 
-        for key, conv_fn in conversion_mapping.items():
-            metadata[key] = conv_fn(metadata[key])
+        metadata_dict = {}
 
-        return metadata
+        for field in dir(metadata):
+            if field.startswith("__"):
+                continue
+            value = getattr(metadata, field)
+            if field in conversion_mapping:
+                value = conversion_mapping[field](value)
+            metadata_dict[field] = value
+
+        return metadata_dict

requirements.txt

@@ -1,5 +1,5 @@
 aiohttp>=3.7.2,<4.0.0
-dbz-python>=0.2.1
+databento-dbn>=0.3.2
 numpy>=1.17.0
 pandas>=1.1.3
 requests>=2.24.0