diff --git a/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py b/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py index 5090088e0df..90a4f53487e 100644 --- a/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py +++ b/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py @@ -207,8 +207,8 @@ if __name__ == "__main__": logger.error("PROVIDER ERROR: WebSocket connection closed") except Exception as e: # pylint: disable=broad-except - msg = f"PROVIDER ERROR: {e.__class__.__name__}: {e}" - logger.error(msg) + ERR = f"PROVIDER ERROR: {e.__class__.__name__}: {e}" + logger.error(ERR) finally: sys.exit(0) 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 2d63f20a9e0..62b94233129 100644 --- a/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py +++ b/openbb_platform/providers/intrinio/openbb_intrinio/utils/websocket_client.py @@ -77,8 +77,8 @@ async def subscribe(symbol, event): elif event == "unsubscribe": client.leave(ticker) except Exception as e: # pylint: disable=broad-except - msg = f"PROVIDER ERROR: {e.__class__.__name__}: {e}" - logger.error(msg) + exc = f"PROVIDER ERROR: {e.__class__.__name__}: {e}" + logger.error(exc) async def read_stdin_and_queue_commands(): @@ -133,8 +133,8 @@ if __name__ == "__main__": logger.error("PROVIDER INFO: WebSocket connection closed") except Exception as e: # pylint: disable=broad-except - msg = f"PROVIDER ERROR: {e.__class__.__name__}: {e}" - logger.error(msg) + EXC = f"PROVIDER ERROR: {e.__class__.__name__}: {e}" + logger.error(EXC) finally: client.disconnect() diff --git a/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py b/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py index b2d8cc0f5ed..305b779da14 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py +++ b/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py @@ -122,7 +122,8 @@ def validate_date(cls, v): v = v / 1e9 # Convert nanoseconds to seconds dt = datetime.fromtimestamp(v, tz=timezone("UTC")) dt = dt.astimezone(timezone("America/New_York")) - return dt + + return dt class PolygonWebSocketQueryParams(WebSocketQueryParams): @@ -352,9 +353,9 @@ class PolygonCryptoTradeWebSocketData(WebSocketData): """Validate the conditions.""" if v is None or isinstance(v, list) and v[0] == 0: return None - elif isinstance(v, list) and v[0] == 1: + if isinstance(v, list) and v[0] == 1: return "sellside" - elif isinstance(v, list) and v[0] == 2: + if isinstance(v, list) and v[0] == 2: return "buyside" return str(v) @@ -1199,7 +1200,7 @@ class PolygonWebSocketData(Data): or MODEL_MAP.get(data.get("type", "")) ) if not model: - return super().__new__(cls) + return super().__new__(cls) # type: ignore return model.model_validate(data) # type: ignore diff --git a/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py b/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py index 891201c2058..c1749b9bb8a 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py +++ b/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py @@ -30,6 +30,7 @@ kwargs = parse_kwargs() CONNECT_KWARGS = kwargs.pop("connect_kwargs", {}) FEED = kwargs.pop("feed", None) ASSET_TYPE = kwargs.pop("asset_type", None) +kwargs["results_file"] = os.path.abspath(kwargs["results_file"]) async def handle_symbol(symbol): @@ -100,16 +101,14 @@ async def login(websocket, api_key): err = f"UnauthorizedError -> {msg.get('message')}" logger.error(err) sys.exit(1) - break if msg.get("status") != "auth_success": err = ( f"UnauthorizedError -> {msg.get('status')} -> {msg.get('message')}" ) logger.error(err) sys.exit(1) - break logger.info("PROVIDER INFO: %s", msg.get("message")) - except Exception as e: + except Exception as e: # pylint: disable=broad-except logger.error("PROVIDER ERROR: %s -> %s", e.__class__.__name__, e.args[0]) sys.exit(1) @@ -124,12 +123,12 @@ async def subscribe(websocket, symbol, event): subscribe_event = f'{{"action":"{event}","params":"{ticker}"}}' try: await websocket.send(subscribe_event) - except Exception as e: + except Exception as e: # pylint: disable=broad-except msg = f"PROVIDER ERROR: {e.__class__.__name__} -> {e}" logger.error(msg) -async def read_stdin(command_queue): +async def read_stdin(): """Read from stdin and queue commands.""" while True: line = await asyncio.get_event_loop().run_in_executor(None, sys.stdin.readline) @@ -155,7 +154,7 @@ async def process_stdin_queue(websocket): await subscribe(websocket, symbol, event) -async def process_message(message, results_path, table_name, limit): +async def process_message(message): """Process the WebSocket message.""" messages = message if isinstance(message, list) else [message] for msg in messages: @@ -168,7 +167,6 @@ async def process_message(message, results_path, table_name, limit): err = f"UnauthorizedError -> {msg.get('message')}" logger.error(err) sys.exit(1) - break logger.info("PROVIDER INFO: %s", msg.get("message")) elif msg and "ev" in msg and "status" not in msg: @@ -183,20 +181,23 @@ async def process_message(message, results_path, table_name, limit): raise e from e if result: - await write_to_db(result, results_path, table_name, limit) + await write_to_db( + result, + kwargs["results_path"], + kwargs["table_name"], + kwargs.get("limit"), + ) else: logger.info("PROVIDER INFO: %s", msg) -async def connect_and_stream(url, symbol, api_key, results_path, table_name, limit): +async def connect_and_stream(): # pylint: disable=too-many-branches, too-many-statements """Connect to the WebSocket and stream data to file.""" handler_task = asyncio.create_task( - queue.process_queue( - lambda message: process_message(message, results_path, table_name, limit) - ) + queue.process_queue(lambda message: process_message(message)) ) - stdin_task = asyncio.create_task(read_stdin(command_queue)) + stdin_task = asyncio.create_task(read_stdin()) try: connect_kwargs = CONNECT_KWARGS.copy() if "ping_timeout" not in connect_kwargs: @@ -205,16 +206,15 @@ async def connect_and_stream(url, symbol, api_key, results_path, table_name, lim connect_kwargs["close_timeout"] = None try: - async with websockets.connect(url, **connect_kwargs) as websocket: - - await login(websocket, api_key) + async with websockets.connect(kwargs["url"], **connect_kwargs) as websocket: + await login(websocket, kwargs["api_key"]) response = await websocket.recv() messages = json.loads(response) - await process_message(messages, results_path, table_name, limit) - await subscribe(websocket, symbol, "subscribe") + await process_message(messages) + await subscribe(websocket, kwargs["symbol"], "subscribe") response = await websocket.recv() messages = json.loads(response) - await process_message(messages, results_path, table_name, limit) + await process_message(messages) while True: cmd_task = asyncio.create_task(process_stdin_queue(websocket)) msg_task = asyncio.create_task(websocket.recv()) @@ -255,14 +255,14 @@ async def connect_and_stream(url, symbol, api_key, results_path, table_name, lim # Attempt to reopen the connection logger.info("PROVIDER INFO: Attempting to reconnect after five seconds.") await asyncio.sleep(5) - await connect_and_stream(url, symbol, api_key, results_path, table_name, limit) + await connect_and_stream() except websockets.WebSocketException as e: msg = f"PROVIDER ERROR: WebSocketException -> {e}" logger.error(msg) sys.exit(1) - except Exception as e: + except Exception as e: # pylint: disable=broad-except msg = f"PROVIDER ERROR: Unexpected error -> {e.__class__.__name__}: {e}" logger.error(msg) sys.exit(1) @@ -284,14 +284,7 @@ if __name__ == "__main__": loop.add_signal_handler(sig, handle_termination_signal, logger) asyncio.run_coroutine_threadsafe( - connect_and_stream( - kwargs["url"], - kwargs["symbol"], - kwargs["api_key"], - os.path.abspath(kwargs["results_file"]), - kwargs["table_name"], - kwargs.get("limit", None), - ), + connect_and_stream(), loop, ) loop.run_forever() @@ -300,8 +293,8 @@ if __name__ == "__main__": logger.error("PROVIDER ERROR: WebSocket connection closed") except Exception as e: # pylint: disable=broad-except - msg = f"PROVIDER ERROR: {e.__class__.__name__} -> {e}" - logger.error(msg) + ERR = f"PROVIDER ERROR: {e.__class__.__name__} -> {e}" + logger.error(ERR) finally: sys.exit(0)