Add WebSocket support

Author: dhruvan2006

requirements.txt

@@ -11,4 +11,6 @@ html5lib>=1.1
 peewee>=3.16.2
 requests_cache>=1.0
 requests_ratelimiter>=0.3.1
-scipy>=1.6.3
+scipy>=1.6.3
+protobuf>=5.29.2
+websockets>=14.1

setup.py

@@ -63,11 +63,17 @@
                       'requests>=2.31', 'multitasking>=0.0.7',
                       'lxml>=4.9.1', 'platformdirs>=2.0.0', 'pytz>=2022.5',
                       'frozendict>=2.3.4', 'peewee>=3.16.2',
-                      'beautifulsoup4>=4.11.1', 'html5lib>=1.1'],
+                      'beautifulsoup4>=4.11.1', 'html5lib>=1.1',
+                      'protobuf>=5.29.2', 'websockets>=14.1'],
     extras_require={
         'nospam': ['requests_cache>=1.0', 'requests_ratelimiter>=0.3.1'],
         'repair': ['scipy>=1.6.3'],
     },
+    # Include protobuf files for websocket support
+    package_data={
+        'yfinance': ['pricing.proto', 'pricing_pb2.py'],
+    },
+    include_package_data=True,
     # Note: Pandas.read_html() needs html5lib & beautifulsoup4
     entry_points={
         'console_scripts': [

tests/test_live.py

@@ -0,0 +1,30 @@
+import unittest
+from unittest.mock import Mock
+
+from yfinance.live import BaseWebSocket
+
+
+class TestWebSocket(unittest.TestCase):
+    def test_decode_message_valid(self):
+        message = ("CgdCVEMtVVNEFYoMuUcYwLCVgIplIgNVU0QqA0NDQzApOAFFPWrEP0iAgOrxvANVx/25R12csrRHZYD8skR9/"
+                   "7i0R7ABgIDq8bwD2AEE4AGAgOrxvAPoAYCA6vG8A/IBA0JUQ4ECAAAAwPrjckGJAgAA2P5ZT3tC")
+
+        ws = BaseWebSocket(Mock())
+        decoded = ws._decode_message(message)
+
+        expected = {'id': 'BTC-USD', 'price': 94745.08, 'time': '1736509140000', 'currency': 'USD', 'exchange': 'CCC',
+                    'quote_type': 41, 'market_hours': 1, 'change_percent': 1.5344921, 'day_volume': '59712028672',
+                    'day_high': 95227.555, 'day_low': 92517.22, 'change': 1431.8906, 'open_price': 92529.99,
+                    'last_size': '59712028672', 'price_hint': '2', 'vol_24hr': '59712028672',
+                    'vol_all_currencies': '59712028672', 'from_currency': 'BTC', 'circulating_supply': 19808172.0,
+                    'market_cap': 1876726640000.0}
+
+        self.assertEqual(expected, decoded)
+
+    def test_decode_message_invalid(self):
+        websocket = BaseWebSocket(Mock())
+        base64_message = "invalid_base64_string"
+        decoded = websocket._decode_message(base64_message)
+        assert "error" in decoded
+        assert "raw_base64" in decoded
+        self.assertEqual(base64_message, decoded["raw_base64"])

yfinance/__init__.py

@@ -24,6 +24,7 @@
 from .ticker import Ticker
 from .tickers import Tickers
 from .multi import download
+from .live import WebSocket, AsyncWebSocket
 from .utils import enable_debug_mode
 from .cache import set_tz_cache_location
 from .domain.sector import Sector
@@ -39,6 +40,6 @@
 import warnings
 warnings.filterwarnings('default', category=DeprecationWarning, module='^yfinance')
 
-__all__ = ['download', 'Market', 'Search', 'Ticker', 'Tickers', 'enable_debug_mode', 'set_tz_cache_location', 'Sector', 'Industry']
+__all__ = ['download', 'Market', 'Search', 'Ticker', 'Tickers', 'enable_debug_mode', 'set_tz_cache_location', 'Sector', 'Industry', 'WebSocket', 'AsyncWebSocket']
 # screener stuff:
 __all__ += ['EquityQuery', 'FundQuery', 'screen', 'PREDEFINED_SCREENER_QUERIES']

yfinance/base.py

@@ -34,6 +34,7 @@
 from . import utils, cache
 from .data import YfData
 from .exceptions import YFEarningsDateMissing, YFRateLimitError
+from .live import WebSocket
 from .scrapers.analysis import Analysis
 from .scrapers.fundamentals import Fundamentals
 from .scrapers.holders import Holders
@@ -76,6 +77,9 @@ def __init__(self, ticker, session=None, proxy=None):
 
         self._fast_info = None
 
+        self._message_handler = None
+        self.ws = None
+
     @utils.log_indent_decorator
     def history(self, *args, **kwargs) -> pd.DataFrame:
         return self._lazy_load_price_history().history(*args, **kwargs)
@@ -690,3 +694,10 @@ def get_funds_data(self, proxy=None) -> Optional[FundsData]:
             self._funds_data = FundsData(self._data, self.ticker)
 
         return self._funds_data
+
+    def live(self, message_handler=None, verbose=True):
+        self._message_handler = message_handler
+
+        self.ws = WebSocket(verbose=verbose)
+        self.ws.subscribe(self.ticker)
+        self.ws.listen(self._message_handler)

yfinance/live.py

@@ -0,0 +1,277 @@
+import asyncio
+import base64
+import json
+from typing import List, Optional, Callable
+
+from websockets.sync.client import connect as sync_connect
+from websockets.asyncio.client import connect as async_connect
+
+from yfinance import utils
+from yfinance.pricing_pb2 import PricingData
+from google.protobuf.json_format import MessageToDict
+
+
+class BaseWebSocket:
+    def __init__(self, url: str = "wss://streamer.finance.yahoo.com/?version=2", verbose=True):
+        self.url = url
+        self.verbose = verbose
+        self.logger = utils.get_yf_logger()
+        self._ws = None
+        self._subscriptions = set()
+        self._subscription_interval = 15  # seconds
+
+    def _decode_message(self, base64_message: str) -> dict:
+        try:
+            decoded_bytes = base64.b64decode(base64_message)
+            pricing_data = PricingData()
+            pricing_data.ParseFromString(decoded_bytes)
+            return MessageToDict(pricing_data, preserving_proto_field_name=True)
+        except Exception as e:
+            self.logger.error("Failed to decode message: %s", e, exc_info=True)
+            print("Failed to decode message: %s", e)
+            return {
+                'error': str(e),
+                'raw_base64': base64_message
+            }
+
+
+class AsyncWebSocket(BaseWebSocket):
+    """
+    Asynchronous WebSocket client for streaming real time pricing data.
+    """
+
+    def __init__(self, url: str = "wss://streamer.finance.yahoo.com/?version=2", verbose=True):
+        """
+        Initialize the AsyncWebSocket client.
+
+        Args:
+            url (str): The WebSocket server URL. Defaults to Yahoo Finance's WebSocket URL.
+            verbose (bool): Flag to enable or disable print statements. Defaults to True.
+        """
+        super().__init__(url, verbose)
+        self._message_handler = None  # Callable to handle messages
+        self._heartbeat_task = None  # Task to send heartbeat subscribe
+
+    async def _connect(self):
+        if self._ws is None:
+            self._ws = await async_connect(self.url)
+            self.logger.info("Connected to WebSocket.")
+            if self.verbose:
+                print("Connected to WebSocket.")
+
+    async def _periodic_subscribe(self):
+        while True:
+            try:
+                await asyncio.sleep(self._subscription_interval)
+
+                if self._subscriptions:
+                    message = {"subscribe": list(self._subscriptions)}
+                    await self._ws.send(json.dumps(message))
+
+                    if self.verbose:
+                        print(f"Heartbeat subscription sent for symbols: {self._subscriptions}")
+            except Exception as e:
+                self.logger.error("Error in heartbeat subscription: %s", e, exc_info=True)
+                if self.verbose:
+                    print(f"Error in heartbeat subscription: {e}")
+                break
+
+    async def subscribe(self, symbols: str | List[str]):
+        """
+        Subscribe to a stock symbol or a list of stock symbols.
+
+        Args:
+            symbols (str | List[str]): Stock symbol(s) to subscribe to.
+        """
+        await self._connect()
+
+        if isinstance(symbols, str):
+            symbols = [symbols]
+
+        self._subscriptions.update(symbols)
+
+        message = {"subscribe": list(self._subscriptions)}
+        await self._ws.send(json.dumps(message))
+
+        # Start heartbeat subscription task
+        if self._heartbeat_task is None:
+            self._heartbeat_task = asyncio.create_task(self._periodic_subscribe())
+
+        self.logger.info(f"Subscribed to symbols: {symbols}")
+        if self.verbose:
+            print(f"Subscribed to symbols: {symbols}")
+
+    async def unsubscribe(self, symbols: str | List[str]):
+        """
+        Unsubscribe from a stock symbol or a list of stock symbols.
+
+        Args:
+            symbols (str | List[str]): Stock symbol(s) to unsubscribe from.
+        """
+        await self._connect()
+
+        if isinstance(symbols, str):
+            symbols = [symbols]
+
+        self._subscriptions.difference_update(symbols)
+
+        message = {"unsubscribe": list(self._subscriptions)}
+        await self._ws.send(json.dumps(message))
+
+        self.logger.info(f"Unsubscribed from symbols: {symbols}")
+        if self.verbose:
+            print(f"Unsubscribed from symbols: {symbols}")
+
+    async def listen(self, message_handler: Optional[Callable[[dict], None]] = None):
+        """
+        Start listening to messages from the WebSocket server.
+
+        Args:
+            message_handler (Optional[Callable[[dict], None]]): Optional function to handle received messages.
+        """
+        await self._connect()
+        self._message_handler = message_handler
+
+        self.logger.info("Listening for messages...")
+        if self.verbose:
+            print("Listening for messages...")
+
+        # Start heartbeat subscription task
+        if self._heartbeat_task is None:
+            self._heartbeat_task = asyncio.create_task(self._periodic_subscribe())
+
+        try:
+            async for message in self._ws:
+                message_json = json.loads(message)
+                encoded_data = message_json.get("message", "")
+                decoded_message = self._decode_message(encoded_data)
+                if self._message_handler:
+                    self._message_handler(decoded_message)
+                else:
+                    print(decoded_message)
+        except (KeyboardInterrupt, asyncio.CancelledError) as e:
+            self.logger.info("WebSocket listening interrupted. Closing connection...")
+            if self.verbose:
+                print("WebSocket listening interrupted. Closing connection...")
+            await self.close()
+        except Exception as e:
+            self.logger.error("Error while listening to messages: %s", e, exc_info=True)
+            if self.verbose:
+                print("Error while listening to messages: %s", e)
+
+    async def close(self):
+        """Close the WebSocket connection."""
+        if self._heartbeat_task:
+            self._heartbeat_task.cancel()
+
+        if self._ws is not None:# and not self._ws.closed:
+            await self._ws.close()
+            self.logger.info("WebSocket connection closed.")
+            if self.verbose:
+                print("WebSocket connection closed.")
+
+
+class WebSocket(BaseWebSocket):
+    """
+    Synchronous WebSocket client for streaming real time pricing data.
+    """
+
+    def __init__(self, url: str = "wss://streamer.finance.yahoo.com/?version=2", verbose=True):
+        """
+        Initialize the WebSocket client.
+
+        Args:
+            url (str): The WebSocket server URL. Defaults to Yahoo Finance's WebSocket URL.
+            verbose (bool): Flag to enable or disable print statements. Defaults to True.
+        """
+        super().__init__(url, verbose)
+
+    def _connect(self):
+        if self._ws is None:
+            self._ws = sync_connect(self.url)
+            self.logger.info("Connected to WebSocket.")
+            if self.verbose:
+                print("Connected to WebSocket.")
+
+    def subscribe(self, symbols: str | List[str]):
+        """
+        Subscribe to a stock symbol or a list of stock symbols.
+
+        Args:
+            symbols (str | List[str]): Stock symbol(s) to subscribe to.
+        """
+        self._connect()
+
+        if isinstance(symbols, str):
+            symbols = [symbols]
+
+        self._subscriptions.update(symbols)
+
+        message = {"subscribe": list(self._subscriptions)}
+        self._ws.send(json.dumps(message))
+
+        self.logger.info(f"Subscribed to symbols: {symbols}")
+        if self.verbose:
+            print(f"Subscribed to symbols: {symbols}")
+
+    def unsubscribe(self, symbols: str | List[str]):
+        """
+        Unsubscribe from a stock symbol or a list of stock symbols.
+
+        Args:
+            symbols (str | List[str]): Stock symbol(s) to unsubscribe from.
+        """
+        self._connect()
+
+        if isinstance(symbols, str):
+            symbols = [symbols]
+
+        self._subscriptions.difference_update(symbols)
+
+        message = {"unsubscribe": list(self._subscriptions)}
+        self._ws.send(json.dumps(message))
+
+        self.logger.info(f"Unsubscribed from symbols: {symbols}")
+        if self.verbose:
+            print(f"Unsubscribed from symbols: {symbols}")
+
+    def listen(self, message_handler: Optional[Callable[[dict], None]] = None):
+        """
+        Start listening to messages from the WebSocket server.
+
+        Args:
+            message_handler (Optional[Callable[[dict], None]]): Optional function to handle received messages.
+        """
+        self._connect()
+
+        self.logger.info("Listening for messages...")
+        if self.verbose:
+            print("Listening for messages...")
+
+        try:
+            while True:
+                message = self._ws.recv()
+                message_json = json.loads(message)
+                encoded_data = message_json.get("message", "")
+                decoded_message = self._decode_message(encoded_data)
+
+                if message_handler:
+                    message_handler(decoded_message)
+                else:
+                    print(decoded_message)
+        except KeyboardInterrupt:
+            if self.verbose:
+                print("Received keyboard interrupt.")
+            self.close()
+        except Exception as e:
+            self.logger.error("Error while listening to messages: %s", e, exc_info=True)
+            if self.verbose:
+                print("Error while listening to messages: %s", e)
+
+    def close(self):
+        """Close the WebSocket connection."""
+        if self._ws is not None:
+            self._ws.close()
+            self.logger.info("WebSocket connection closed.")
+            if self.verbose:
+                print("WebSocket connection closed.")