From 53cd9987b79c2a3285bf4f8ba77b881b46a19de3 Mon Sep 17 00:00:00 2001 From: Danglewood <85772166+deeleeramone@users.noreply.github.com> Date: Mon, 9 Dec 2024 17:22:16 -0800 Subject: [PATCH] some more intrinio.. --- .../openbb_intrinio/models/websocket_connection.py | 2 +- .../intrinio/openbb_intrinio/utils/stocks_client.py | 2 +- .../openbb_intrinio/utils/websocket_client.py | 12 ++++-------- 3 files changed, 6 insertions(+), 10 deletions(-) diff --git a/openbb_platform/providers/intrinio/openbb_intrinio/models/websocket_connection.py b/openbb_platform/providers/intrinio/openbb_intrinio/models/websocket_connection.py index d10d5b1c0c8..b3b2c853337 100644 --- a/openbb_platform/providers/intrinio/openbb_intrinio/models/websocket_connection.py +++ b/openbb_platform/providers/intrinio/openbb_intrinio/models/websocket_connection.py @@ -62,7 +62,7 @@ class IntrinioWebSocketData(WebSocketData): default=None, description="The exchange of the data.", ) - type: Literal["quote", "trade"] = Field( + type: Literal["bid", "ask", "trade"] = Field( description="The type of data.", ) price: Optional[float] = Field( diff --git a/openbb_platform/providers/intrinio/openbb_intrinio/utils/stocks_client.py b/openbb_platform/providers/intrinio/openbb_intrinio/utils/stocks_client.py index 318f0bf6862..aeb2f31bf60 100644 --- a/openbb_platform/providers/intrinio/openbb_intrinio/utils/stocks_client.py +++ b/openbb_platform/providers/intrinio/openbb_intrinio/utils/stocks_client.py @@ -112,7 +112,7 @@ class Trade: self.timestamp = timestamp self.subprovider = subprovider self.market_center = market_center - self.condition = str(condition).strip().replace(" ", "") + self.condition = condition def __str__(self): """Return string representation of the trade.""" diff --git a/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py b/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py index 49764c6dcd5..2d63f20a9e0 100644 --- a/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py +++ b/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py @@ -9,7 +9,7 @@ import sys from openbb_core.provider.utils.helpers import run_async from openbb_intrinio.models.websocket_connection import IntrinioWebSocketData -from openbb_intrinio.utils.stocks_client import IntrinioRealtimeClient, Quote, Trade +from openbb_intrinio.utils.stocks_client import IntrinioRealtimeClient, Trade from openbb_websockets.helpers import ( MessageQueue, get_logger, @@ -31,17 +31,13 @@ async def process_message(message): result: dict = {} message = json.loads(message) if isinstance(message, str) else message is_trade = isinstance(message, Trade) - is_quote = isinstance(message, Quote) if hasattr(message, "__dict__"): message = message.__dict__ - if is_trade or is_quote: + if is_trade: message["type"] = "trade" if is_trade else "quote" try: - result = IntrinioWebSocketData.model_validate(message).model_dump_json( - exclude_none=True, exclude_unset=True - ) - result = message + result = IntrinioWebSocketData.model_validate(message).model_dump_json() except ValidationError as e: try: handle_validation_error(logger, e) @@ -49,7 +45,7 @@ async def process_message(message): raise e from e if result: await write_to_db( - message, kwargs["results_file"], kwargs["table_name"], kwargs["limit"] + result, kwargs["results_file"], kwargs["table_name"], kwargs["limit"] )