From 3ca267fd5795459bf2b5a73b9186262a1f48942c Mon Sep 17 00:00:00 2001 From: Danglewood <85772166+deeleeramone@users.noreply.github.com> Date: Fri, 15 Nov 2024 11:36:21 -0800 Subject: [PATCH] docstring things --- .../extensions/websockets/README.md | 80 ++++++++++++++++++- .../websockets/openbb_websockets/client.py | 75 +++++++++-------- .../websockets/openbb_websockets/helpers.py | 2 +- .../openbb_websockets/websockets_router.py | 5 +- 4 files changed, 119 insertions(+), 43 deletions(-) diff --git a/openbb_platform/extensions/websockets/README.md b/openbb_platform/extensions/websockets/README.md index b98a11e2aaf..ee9c120a3ce 100644 --- a/openbb_platform/extensions/websockets/README.md +++ b/openbb_platform/extensions/websockets/README.md @@ -454,12 +454,11 @@ PROVIDER INFO: unsubscribed to: C.XAU/USD ## Development - ### Provider Interface Providers can be added to the `create_connection` endpoint by following a slightly modified pattern. This section outlines the adaptations, but does not contain any code for actually connecting to the provider's websocket. -For details on that part, go to [websocket_client](README.md###websocket_client) section below. +For details on that part, go to [websocket_client](README.md#websocket_client) section below. Here, the Fetcher is used to start the provider client module (in a separate file) and return the client to the router, where it is intercepted and kept alive. @@ -621,6 +620,83 @@ class FmpWebSocketConnection(WebSocketConnection): """FMP WebSocket connection model.""" ``` +#### WebSocketClient + +The `WebSocketClient` is the main class responsible for bidrectional communication between the provider, broadcast, and user. It handles the child processes and can be used as a standalone class. Pasted below is the docstring for the class. + +It can be imported to use as a standalone class, and this instance is the 'client' field in `WebSocketConnection` + +```python +from openbb_websockets.client import WebSocketClient +``` + +```console + Parameters + ---------- + name : str + Name to assign the WebSocket connection. Used to identify and manage multiple instances. + module : str + The Python module for the provider websocket_client module. Runs in a separate thread. + Example: 'openbb_fmp.utils.websocket_client'. Pass additional keyword arguments by including kwargs. + symbol : Optional[str] + The symbol(s) requested to subscribe. Enter multiple symbols separated by commas without spaces. + limit : Optional[int] + The limit of records to hold in memory. Once the limit is reached, the oldest records are removed. + Default is 1000. Set to None to keep all records. + results_file : Optional[str] + Absolute path to the file for continuous writing. By default, a temporary file is created. + table_name : Optional[str] + SQL table name to store serialized data messages. By default, 'records'. + save_results : bool + Whether to persist the results after the main Python session ends. Default is False. + data_model : Optional[Data] + Pydantic data model to validate the results before storing them in the database. + Also used to deserialize the results from the database. + auth_token : Optional[str] + The authentication token to use for the WebSocket connection. Default is None. + Only used for API and Python application endpoints. + logger : Optional[logging.Logger] + The pre-configured logger instance to use for this connection. By default, a new logger is created. + kwargs : Optional[dict] + Additional keyword arguments to pass to the target provider module. Keywords and values must not contain spaces. + To pass items to 'websocket.connect()', include them in the 'kwargs' dictionary as, + {'connect_kwargs': {'key': 'value'}}. + + Properties + ---------- + symbol : str + Symbol(s) requested to subscribe. + module : str + Path to the provider connection script. + is_running : bool + Check if the provider connection process is running. + is_broadcasting : bool + Check if the broadcast server process is running. + broadcast_address : str + URI address for the results broadcast server. + results : list[Data] + All stored results from the provider's WebSocket stream. + Results are stored in a SQLite database as a serialized JSON string, this property deserializes the results. + Clear the results by deleting the property. e.g., del client.results + + Methods + ------- + connect + Connect to the provider WebSocket stream. + disconnect + Disconnect from the provider WebSocket. + subscribe + Subscribe to a new symbol or list of symbols. + unsubscribe + Unsubscribe from a symbol or list of symbols. + start_broadcasting + Start the broadcast server to stream results over a network connection. + stop_broadcasting + Stop the broadcast server and disconnect all listening clients. + send_message + Send a message to the WebSocket process. Messages can be sent to "provider" or "broadcast" targets. +``` + #### WebSocketFetcher This is where things diverge slightly. Instead of returning `FmpWebSocketData`, it gets passed to the client connection for validating records as they are received. What gets returned by the Fetcher is the `WebSocketConnection`. diff --git a/openbb_platform/extensions/websockets/openbb_websockets/client.py b/openbb_platform/extensions/websockets/openbb_websockets/client.py index 72a720b3ce5..1aa7178ce50 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/client.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/client.py @@ -3,7 +3,7 @@ # pylint: disable=too-many-statements # flake8: noqa: PLR0915 import logging -from typing import TYPE_CHECKING, Literal, Optional, Union +from typing import TYPE_CHECKING, Any, Literal, Optional, Union if TYPE_CHECKING: from openbb_core.provider.abstract.data import Data @@ -17,13 +17,13 @@ class WebSocketClient: name : str Name to assign the WebSocket connection. Used to identify and manage multiple instances. module : str - The Python module for the provider server connection. Runs in a separate thread. - Example: 'openbb_fmp.websockets.server'. Pass additional keyword arguments by including kwargs. + The Python module for the provider websocket_client module. Runs in a separate thread. + Example: 'openbb_fmp.utils.websocket_client'. Pass additional keyword arguments by including kwargs. symbol : Optional[str] The symbol(s) requested to subscribe. Enter multiple symbols separated by commas without spaces. limit : Optional[int] The limit of records to hold in memory. Once the limit is reached, the oldest records are removed. - Default is 300. Set to None to keep all records. + Default is 1000. Set to None to keep all records. results_file : Optional[str] Absolute path to the file for continuous writing. By default, a temporary file is created. table_name : Optional[str] @@ -55,12 +55,10 @@ class WebSocketClient: Check if the broadcast server process is running. broadcast_address : str URI address for the results broadcast server. - results : list - All stored results from the provider's WebSocket stream. The results are stored in a SQLite database. - Set the 'limit' property to cap the number of stored records. + results : list[Data] + All stored results from the provider's WebSocket stream. + Results are stored in a SQLite database as a serialized JSON string, this property deserializes the results. Clear the results by deleting the property. e.g., del client.results - transformed_results : list - Deserialize the records from the results file using the provided data model, if available. Methods ------- @@ -75,9 +73,9 @@ class WebSocketClient: start_broadcasting Start the broadcast server to stream results over a network connection. stop_broadcasting - Stop the broadcast server and disconnect all reading clients. + Stop the broadcast server and disconnect all listening clients. send_message - Send a message to the WebSocket process. + Send a message to the WebSocket process. Messages can be sent to "provider" or "broadcast" targets. """ def __init__( # noqa: PLR0913 @@ -85,7 +83,7 @@ class WebSocketClient: name: str, module: str, symbol: Optional[str] = None, - limit: Optional[int] = 300, + limit: Optional[int] = 1000, results_file: Optional[str] = None, table_name: Optional[str] = None, save_results: bool = False, @@ -103,6 +101,7 @@ class WebSocketClient: from aiosqlite import DatabaseError from queue import Queue from pathlib import Path + from openbb_core.app.model.abstract.error import OpenBBError from openbb_websockets.helpers import get_logger self.name = name @@ -119,20 +118,20 @@ class WebSocketClient: else None ) - self._process = None - self._psutil_process = None - self._thread = None - self._log_thread = None - self._provider_message_queue = Queue() - self._stop_log_thread_event = threading.Event() - self._stop_broadcasting_event = threading.Event() - self._broadcast_address = None - self._broadcast_process = None - self._psutil_broadcast_process = None - self._broadcast_thread = None - self._broadcast_log_thread = None - self._broadcast_message_queue = Queue() - self._exception = None + self._process: Any = None + self._psutil_process: Any = None + self._thread: Any = None + self._log_thread: Any = None + self._provider_message_queue: Queue = Queue() + self._stop_log_thread_event: threading.Event = threading.Event() + self._stop_broadcasting_event: threading.Event = threading.Event() + self._broadcast_address: Any = None + self._broadcast_process: Any = None + self._psutil_broadcast_process: Any = None + self._broadcast_thread: Any = None + self._broadcast_log_thread: Any = None + self._broadcast_message_queue: Queue = Queue() + self._exception: Any = None if not results_file: with tempfile.NamedTemporaryFile(delete=False) as temp_file: @@ -150,7 +149,9 @@ class WebSocketClient: try: self._setup_database() except DatabaseError as e: - self.logger.error("Error setting up the SQLite database and table: %s", e) + msg = f"Unexpected error setting up the SQLite database and table -> {e.__class___.__name__}: {e.__str__()}" + self.logger.error(msg) + self._exception = OpenBBError(msg) def _atexit(self) -> None: """Clean up the WebSocket client processes at exit.""" @@ -164,8 +165,8 @@ class WebSocketClient: if self.is_broadcasting: self.stop_broadcasting() if self.save_results: - self.logger.info("Websocket results saved to, %s\n", self.results_file) - if os.path.exists(self.results_file): + self.logger.info("Websocket results saved to, %s\n", str(self.results_path)) + if os.path.exists(self.results_file) and not self.save_results: os.remove(self.results_file) def _setup_database(self) -> None: @@ -476,7 +477,7 @@ class WebSocketClient: return False @property - def results(self) -> Union[list[dict], None]: + def results(self) -> Union[list[dict], list["Data"], None]: """Retrieve the deserialized results from the results file.""" # pylint: disable=import-outside-toplevel import json # noqa @@ -489,7 +490,12 @@ class WebSocketClient: cursor = conn.execute(f"SELECT * FROM {self.table_name}") # noqa for row in cursor: index, message = row - output.append(json.loads(json.loads(message))) + if self.data_model: + output.append( + self.data_model.model_validate_json(json.loads(message)) + ) + else: + output.append(json.loads(json.loads(message))) if output: return output @@ -665,13 +671,6 @@ class WebSocketClient: self._stop_broadcasting_event.clear() return - @property - def transformed_results(self) -> list["Data"]: - """Model validated records from the results file.""" - if not self.data_model: - raise NotImplementedError("No model provided to transform the results.") - return [self.data_model.model_validate(d) for d in self.results] - def __repr__(self): """Return the WebSocketClient representation.""" return ( diff --git a/openbb_platform/extensions/websockets/openbb_websockets/helpers.py b/openbb_platform/extensions/websockets/openbb_websockets/helpers.py index d8cedfcba97..4dc8b144404 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/helpers.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/helpers.py @@ -41,7 +41,7 @@ def get_logger(name, level=logging.INFO): def handle_validation_error(logger: logging.Logger, error: ValidationError): """Log and raise a Pydantic ValidationError from a provider connection.""" - err = f"{error.__class__.__name__} -> {error.title}: {str(error.json())}" + err = f"{error.__class__.__name__} -> {error.title}: {error.json()}" logger.error(err) raise error from error diff --git a/openbb_platform/extensions/websockets/openbb_websockets/websockets_router.py b/openbb_platform/extensions/websockets/openbb_websockets/websockets_router.py index 5445e0e1b6a..cd2a9c28034 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/websockets_router.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/websockets_router.py @@ -18,6 +18,7 @@ from openbb_core.app.provider_interface import ( from openbb_core.app.query import Query from openbb_core.app.router import Router from openbb_core.provider.utils.errors import EmptyDataError, UnauthorizedError +from pydantic import ValidationError from openbb_websockets.helpers import ( StdOutSink, @@ -101,9 +102,9 @@ async def get_results(name: str, auth_token: Optional[str] = None) -> OBBject: if not client.results: raise EmptyDataError(f"No results recorded for client {name}.") try: - return OBBject(results=client.transformed_results) - except NotImplementedError: return OBBject(results=client.results) + except ValidationError as e: + raise OpenBBError(e) from e @router.command(