ADD: Add Live client symbology mapping

Author: nmacholl

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## 0.15.0 - TBD
+
+### Enhancements
+- Added `symbology_map` property to `Live` client
+
 ## 0.14.1 - 2023-06-16
 
 - Fixed issue where `DBNStore.to_df()` would raise an exception if no records were present

databento/live/client.py

@@ -88,7 +88,8 @@ def __init__(
 
         self._dbn_queue: DBNQueue = DBNQueue(maxsize=DEFAULT_QUEUE_SIZE)
         self._metadata: SessionMetadata = SessionMetadata()
-        self._user_callbacks: list[UserCallback] = []
+        self._symbology_map: dict[int, str | int] = {}
+        self._user_callbacks: list[UserCallback] = [self._map_symbol]
         self._user_streams: list[IO[bytes]] = []
 
         def factory() -> _SessionProtocol:
@@ -237,6 +238,23 @@ def port(self) -> int:
         """
         return self._port
 
+    @property
+    def symbology_map(self) -> dict[int, str | int]:
+        """
+        Return the symbology map for this client session. A symbol mapping is
+        added when the client receives a SymbolMappingMsg.
+
+        This can be used to transform an `instrument_id` in a DBN record
+        to the input symbology.
+
+        Returns
+        -------
+        dict[int, str | int]
+            A mapping of the exchange's instrument_id to the subscription symbology.
+
+        """
+        return self._symbology_map
+
     @property
     def ts_out(self) -> bool:
         """
@@ -548,3 +566,11 @@ async def _shutdown(self) -> None:
         if self._session is None:
             return
         await self._session.wait_for_close()
+        self._symbology_map.clear()
+
+    def _map_symbol(self, record: DBNRecord) -> None:
+        if isinstance(record, databento_dbn.SymbolMappingMsg):
+            out_symbol = record.stype_out_symbol
+            instrument_id = record.instrument_id
+            self._symbology_map[instrument_id] = record.stype_out_symbol
+            logger.info("added symbology mapping %s to %d", out_symbol, instrument_id)

databento/live/protocol.py

@@ -312,8 +312,22 @@ def _process_dbn(self, data: bytes) -> None:
                 logger.debug("dispatching %s", type(record).__name__)
                 if isinstance(record, databento_dbn.Metadata):
                     self.received_metadata(record)
-                else:
-                    self.received_record(record)
+                    continue
+
+                if isinstance(record, databento_dbn.ErrorMsg):
+                    logger.error(
+                        "gateway error: %s",
+                        record.err,
+                    )
+                if isinstance(record, databento_dbn.SystemMsg):
+                    if record.is_heartbeat:
+                        logger.debug("gateway heartbeat")
+                    else:
+                        logger.info(
+                            "gateway message: %s",
+                            record.msg,
+                        )
+                self.received_record(record)
 
     def _process_gateway(self, data: bytes) -> None:
         try:

databento/live/session.py

@@ -470,15 +470,15 @@ async def _connect_task(
                 ),
                 timeout=CONNECT_TIMEOUT_SECONDS,
             )
-        except asyncio.TimeoutError:
+        except asyncio.TimeoutError as exc:
             raise BentoError(
                 f"Connection to {gateway}:{port} timed out after "
                 f"{CONNECT_TIMEOUT_SECONDS} second(s).",
-            )
-        except OSError:
+            ) from exc
+        except OSError as exc:
             raise BentoError(
                 f"Connection to {gateway}:{port} failed.",
-            )
+            ) from exc
 
         logger.debug(
             "connected to %s:%d",
@@ -491,13 +491,13 @@ async def _connect_task(
                 protocol.authenticated,
                 timeout=AUTH_TIMEOUT_SECONDS,
             )
-        except asyncio.TimeoutError:
+        except asyncio.TimeoutError as exc:
             raise BentoError(
                 f"Authentication with {gateway}:{port} timed out after "
                 f"{AUTH_TIMEOUT_SECONDS} second(s).",
-            )
+            ) from exc
         except ValueError as exc:
-            raise BentoError(f"User authentication failed: {str(exc)}")
+            raise BentoError(f"User authentication failed: {str(exc)}") from exc
 
         logger.info(
             "authentication with remote gateway completed",

databento/version.py

@@ -1 +1 @@
-__version__ = "0.14.1"
+__version__ = "0.15.0"

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "databento"
-version = "0.14.1"
+version = "0.15.0"
 description = "Official Python client library for Databento"
 authors = [
     "Databento <[email protected]>",

tests/test_live_client.py

@@ -496,7 +496,7 @@ def callback(_: object) -> None:
         pass
 
     live_client.add_callback(callback)
-    assert live_client._user_callbacks == [callback]
+    assert callback in live_client._user_callbacks
     assert live_client._user_streams == []
 
 
@@ -509,7 +509,6 @@ def test_live_add_stream(
     stream = BytesIO()
 
     live_client.add_stream(stream)
-    assert live_client._user_callbacks == []
     assert live_client._user_streams == [stream]
 
 
@@ -581,7 +580,9 @@ async def test_live_async_iteration_backpressure(
         symbols="TEST",
     )
 
-    monkeypatch.setattr(live_client._session._transport, "pause_reading", pause_mock:=MagicMock())
+    monkeypatch.setattr(
+        live_client._session._transport, "pause_reading", pause_mock := MagicMock(),
+    )
 
     live_client.start()
     it = live_client.__iter__()
@@ -618,7 +619,9 @@ async def test_live_async_iteration_dropped(
         symbols="TEST",
     )
 
-    monkeypatch.setattr(live_client._session._transport, "pause_reading", pause_mock:=MagicMock())
+    monkeypatch.setattr(
+        live_client._session._transport, "pause_reading", pause_mock := MagicMock(),
+    )
 
     live_client.start()
     it = live_client.__iter__()
@@ -630,6 +633,7 @@ async def test_live_async_iteration_dropped(
     assert len(records) == 1
     assert live_client._dbn_queue.empty()
 
+
 @pytest.mark.skipif(platform.system() == "Windows", reason="flaky on windows runner")
 async def test_live_async_iteration_stop(
     live_client: client.Live,