ADD: Add Python client set_symbology_json

Author: cjdsellers

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## 0.31.0 - TBD
+
+#### Enhancements
+- Added `DBNStore.insert_symbology_json` convenience method for adding symbology data from a JSON dict or file path
+
 ## 0.30.0 - 2024-02-22
 
 #### Enhancements

databento/common/dbnstore.py

@@ -6,6 +6,7 @@
 import logging
 from collections.abc import Generator
 from collections.abc import Iterator
+from collections.abc import Mapping
 from io import BytesIO
 from os import PathLike
 from pathlib import Path
@@ -18,6 +19,7 @@
     Final,
     Literal,
     Protocol,
+    TextIO,
     overload,
 )
 
@@ -108,20 +110,16 @@ class DataSource(abc.ABC):
     Abstract base class for backing DBNStore instances with data.
     """
 
-    def __init__(self, source: object) -> None:
-        ...
+    def __init__(self, source: object) -> None: ...
 
     @property
-    def name(self) -> str:
-        ...
+    def name(self) -> str: ...
 
     @property
-    def nbytes(self) -> int:
-        ...
+    def nbytes(self) -> int: ...
 
     @property
-    def reader(self) -> IO[bytes]:
-        ...
+    def reader(self) -> IO[bytes]: ...
 
 
 class FileDataSource(DataSource):
@@ -675,6 +673,27 @@ def from_bytes(cls, data: BytesIO | bytes | IO[bytes]) -> DBNStore:
         """
         return cls(MemoryDataSource(data))
 
+    def insert_symbology_json(
+        self,
+        json_data: str | Mapping[str, Any] | TextIO,
+        clear_existing: bool = True,
+    ) -> None:
+        """
+        Insert the given JSON data obtained from the `symbology.resolve`
+        endpoint or a `symbology.json` file.
+
+        Parameters
+        ----------
+        json_data : str | Mapping[str, Any] | TextIO
+            The JSON data to insert.
+        clear_existing : bool, default True
+            If existing symbology data should be cleared from the internal mappings.
+
+        """
+        if clear_existing:
+            self._instrument_map.clear()
+        self._instrument_map.insert_json(json_data)
+
     def replay(self, callback: Callable[[Any], None]) -> None:
         """
         Replay data by passing records sequentially to the given callback.
@@ -834,8 +853,7 @@ def to_df(
         schema: Schema | str | None = ...,
         tz: pytz.BaseTzInfo | str = ...,
         count: None = ...,
-    ) -> pd.DataFrame:
-        ...
+    ) -> pd.DataFrame: ...
 
     @overload
     def to_df(
@@ -846,16 +864,17 @@ def to_df(
         schema: Schema | str | None = ...,
         tz: pytz.BaseTzInfo | str = ...,
         count: int = ...,
-    ) -> DataFrameIterator:
-        ...
+    ) -> DataFrameIterator: ...
 
     def to_df(
         self,
         price_type: Literal["fixed", "float", "decimal"] = "float",
         pretty_ts: bool = True,
         map_symbols: bool = True,
         schema: Schema | str | None = None,
-        tz: pytz.BaseTzInfo | str | Default[pytz.BaseTzInfo] = Default[pytz.BaseTzInfo](pytz.UTC),
+        tz: pytz.BaseTzInfo | str | Default[pytz.BaseTzInfo] = Default[pytz.BaseTzInfo](
+            pytz.UTC,
+        ),
         count: int | None = None,
     ) -> pd.DataFrame | DataFrameIterator:
         """
@@ -903,7 +922,9 @@ def to_df(
         if isinstance(tz, Default):
             tz = tz.value  # consume default
         elif not pretty_ts:
-            raise ValueError("A timezone was specified when `pretty_ts` is `False`. Did you mean to set `pretty_ts=True`?")
+            raise ValueError(
+                "A timezone was specified when `pretty_ts` is `False`. Did you mean to set `pretty_ts=True`?",
+            )
 
         if not isinstance(tz, pytz.BaseTzInfo):
             tz = pytz.timezone(tz)
@@ -1096,16 +1117,14 @@ def to_ndarray(  # type: ignore [misc]
         self,
         schema: Schema | str | None = ...,
         count: None = ...,
-    ) -> np.ndarray[Any, Any]:
-        ...
+    ) -> np.ndarray[Any, Any]: ...
 
     @overload
     def to_ndarray(
         self,
         schema: Schema | str | None = ...,
         count: int = ...,
-    ) -> NDArrayIterator:
-        ...
+    ) -> NDArrayIterator: ...
 
     def to_ndarray(
         self,
@@ -1242,12 +1261,10 @@ def _schema_struct_map(self) -> dict[Schema, type[DBNRecord]]:
 
 class NDArrayIterator(Protocol):
     @abc.abstractmethod
-    def __iter__(self) -> NDArrayIterator:
-        ...
+    def __iter__(self) -> NDArrayIterator: ...
 
     @abc.abstractmethod
-    def __next__(self) -> np.ndarray[Any, Any]:
-        ...
+    def __next__(self) -> np.ndarray[Any, Any]: ...
 
 
 class NDArrayStreamIterator(NDArrayIterator):

pyproject.toml

@@ -57,6 +57,9 @@ types-pytz = "^2024.1.0.20240203"
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
+[tool.black]
+line_length = 100
+
 [tool.mypy]
 python_version = 3.8
 disallow_untyped_defs = true

tests/conftest.py

@@ -1,6 +1,7 @@
 """
 Pytest fixtures.
 """
+
 from __future__ import annotations
 
 import asyncio
@@ -219,6 +220,7 @@ def fixture_thread_loop() -> Generator[asyncio.AbstractEventLoop, None, None]:
     yield loop
     loop.stop()
 
+
 @pytest.fixture(name="mock_live_server")
 def fixture_mock_live_server(
     thread_loop: asyncio.AbstractEventLoop,
@@ -266,6 +268,7 @@ def fixture_mock_live_server(
             loop=thread_loop,
         ).result()
 
+
 @pytest.fixture(name="historical_client")
 def fixture_historical_client(
     test_api_key: str,

tests/test_common_symbology.py

@@ -961,3 +961,38 @@ def test_instrument_map_symbols_json(
     # Assert
     assert outfile == written_path
     assert outfile.read_text() == expected_path.read_text()
+
+
+def test_insert_symbology_json_mismatched_stypes(
+    test_data_path: Callable[[Dataset, Schema], pathlib.Path],
+) -> None:
+    """
+    Test setting JSON symbology data.
+    """
+    # Arrange
+    store = DBNStore.from_file(test_data_path(Dataset.XNAS_ITCH, Schema.TRADES))
+
+    result = {
+        "NVDA": [
+            {
+                "d0": store.start.date().isoformat(),
+                "d1": store.end.date().isoformat(),
+                "s": "6155",
+            },
+        ],
+    }
+    sym_resp = create_symbology_response(
+        result=result,
+        symbols=store.symbols,
+        stype_in=SType.RAW_SYMBOL,
+        stype_out=SType.INSTRUMENT_ID,
+        start_date=store.start.date(),
+        end_date=store.end.date(),
+    )
+
+    # Act
+    store.insert_symbology_json(sym_resp)
+
+    # Assert
+    assert store.to_df().iloc[0]["symbol"] == "NVDA"
+    assert store.to_df().iloc[0]["instrument_id"] == 6155