goldyard/yfinance-fork
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add WebSocket support

Browse Source
This commit is contained in:
Dhruvan Gnanadhandayuthapani
2025-01-18 14:44:53 +01:00
parent 290bf9b7cf
commit fc5d29558b
14 changed files with 499 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1
View File
@@ -40,6 +40,7 @@
- `Tickers`: multiple tickers' data
- `download`: download market data for multiple tickers
- `Market`: get information about a market
- `WebSocket` and `AsyncWebSocket`: live streaming data
- `Search`: quotes and news from search
- `Sector` and `Industry`: sector and industry information
- `EquityQuery` and `Screener`: build query to screen market
doc/source/reference/examples/live_async.py
+23
View File
@@ -0,0 +1,23 @@
import asyncio
import yfinance as yf
# define your message callback
def message_handler(message):
print("Received message:", message)
async def main():
# =======================
# With Context Manager
# =======================
async with yf.AsyncWebSocket() as ws:
await ws.subscribe(["AAPL", "BTC-USD"])
await ws.listen()
# =======================
# Without Context Manager
# =======================
ws = yf.AsyncWebSocket()
await ws.subscribe(["AAPL", "BTC-USD"])
await ws.listen()
asyncio.run(main())
doc/source/reference/examples/live_sync.py
+19
View File
@@ -0,0 +1,19 @@
import yfinance as yf
# define your message callback
def message_handler(message):
print("Received message:", message)
# =======================
# With Context Manager
# =======================
with yf.WebSocket() as ws:
ws.subscribe(["AAPL", "BTC-USD"])
ws.listen(message_handler)
# =======================
# Without Context Manager
# =======================
ws = yf.WebSocket()
ws.subscribe(["AAPL", "BTC-USD"])
ws.listen(message_handler)
doc/source/reference/examples/ticker.py
+3
View File
@@ -20,3 +20,6 @@ dat.info
# analysis
dat.analyst_price_targets
# websocket
dat.live()
doc/source/reference/examples/tickers.py
+4 -1
View File
@@ -5,4 +5,7 @@ tickers = yf.Tickers('msft aapl goog')
# access each ticker using (example)
tickers.tickers['MSFT'].info
tickers.tickers['AAPL'].history(period="1mo")
tickers.tickers['GOOG'].actions
tickers.tickers['GOOG'].actions
# websocket
tickers.live()
doc/source/reference/index.rst
+3
View File
@@ -19,6 +19,8 @@ The following are the publicly available classes, and functions exposed by the `
- :attr:`download <yfinance.download>`: Function to download market data for multiple tickers.
- :attr:`Search <yfinance.Search>`: Class for accessing search results.
- :attr:`Lookup <yfinance.Lookup>`: Class for looking up tickers.
- :class:`WebSocket <yfinance.WebSocket>`: Class for synchronously streaming live market data.
- :class:`AsyncWebSocket <yfinance.AsyncWebSocket>`: Class for asynchronously streaming live market data.
- :attr:`Sector <yfinance.Sector>`: Domain class for accessing sector information.
- :attr:`Industry <yfinance.Industry>`: Domain class for accessing industry information.
- :attr:`Market <yfinance.Market>`: Class for accessing market status & summary.
@@ -41,6 +43,7 @@ The following are the publicly available classes, and functions exposed by the `
yfinance.market
yfinance.search
yfinance.lookup
yfinance.websocket
yfinance.sector_industry
yfinance.screener
yfinance.functions
doc/source/reference/yfinance.websocket.rst
+46
View File
@@ -0,0 +1,46 @@
=====================
WebSocket
=====================
.. currentmodule:: yfinance
The `WebSocket` module allows you to stream live price data from Yahoo Finance using both synchronous and asynchronous clients.
Classes
------------
.. autosummary::
:toctree: api/
WebSocket
AsyncWebSocket
Synchronous WebSocket
----------------------
The `WebSocket` class provides a synchronous interface for subscribing to price updates.
Sample Code:
.. literalinclude:: examples/live_sync.py
:language: python
Asynchronous WebSocket
-----------------------
The `AsyncWebSocket` class provides an asynchronous interface for subscribing to price updates.
Sample Code:
.. literalinclude:: examples/live_async.py
:language: python
.. note::
If you're running asynchronous code in a Jupyter notebook, you may encounter issues with event loops. To resolve this, you need to import and apply `nest_asyncio` to allow nested event loops.
Add the following code before running asynchronous operations:
.. code-block:: python
import nest_asyncio
nest_asyncio.apply()
requirements.txt
+3 -1
View File
@@ -9,4 +9,6 @@ beautifulsoup4>=4.11.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.0,<6
websockets>=11.0
setup.py
+7 -1
View File
@@ -63,11 +63,17 @@ setup(
'requests>=2.31', 'multitasking>=0.0.7',
'platformdirs>=2.0.0', 'pytz>=2022.5',
'frozendict>=2.3.4', 'peewee>=3.16.2',
'beautifulsoup4>=4.11.1'],
'beautifulsoup4>=4.11.1',
'protobuf>=5.29.0,<6', 'websockets>=11.0'],
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
+30
View File
@@ -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
+2 -1
View File
@@ -25,6 +25,7 @@ from .lookup import Lookup
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
@@ -41,7 +42,7 @@ __author__ = "Ran Aroussi"
import warnings
warnings.filterwarnings('default', category=DeprecationWarning, module='^yfinance')
__all__ = ['download', 'Market', 'Search', 'Lookup', 'Ticker', 'Tickers', 'enable_debug_mode', 'set_tz_cache_location', 'Sector', 'Industry']
__all__ = ['download', 'Market', 'Search', 'Lookup', '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
+11
View File
@@ -33,6 +33,7 @@ import requests
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
@@ -86,6 +87,9 @@ class TickerBase:
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)
@@ -792,3 +796,10 @@ class TickerBase:
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
+335
View File
@@ -0,0 +1,335 @@
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)
if self.verbose:
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):
try:
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.")
except Exception as e:
self.logger.error("Failed to connect to WebSocket: %s", e, exc_info=True)
if self.verbose:
print(f"Failed to connect to WebSocket: {e}")
self._ws = None
raise
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": symbols}
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=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())
while True:
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:
try:
if asyncio.iscoroutinefunction(self._message_handler):
await self._message_handler(decoded_message)
else:
self._message_handler(decoded_message)
except Exception as handler_exception:
self.logger.error("Error in message handler: %s", handler_exception, exc_info=True)
if self.verbose:
print("Error in message handler:", handler_exception)
else:
print(decoded_message)
except (KeyboardInterrupt, asyncio.CancelledError):
self.logger.info("WebSocket listening interrupted. Closing connection...")
if self.verbose:
print("WebSocket listening interrupted. Closing connection...")
await self.close()
break
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)
# Attempt to reconnect if connection drops
self.logger.info("Attempting to reconnect...")
if self.verbose:
print("Attempting to reconnect...")
await asyncio.sleep(3) # backoff
await self._connect()
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.")
async def __aenter__(self):
await self._connect()
return self
async def __aexit__(self, exc_type, exc_value, traceback):
await self.close()
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):
try:
if self._ws is None:
self._ws = sync_connect(self.url)
self.logger.info("Connected to WebSocket.")
if self.verbose:
print("Connected to WebSocket.")
except Exception as e:
self.logger.error("Failed to connect to WebSocket: %s", e, exc_info=True)
if self.verbose:
print(f"Failed to connect to WebSocket: {e}")
self._ws = None
raise
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": symbols}
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...")
while True:
try:
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:
try:
message_handler(decoded_message)
except Exception as handler_exception:
self.logger.error("Error in message handler: %s", handler_exception, exc_info=True)
if self.verbose:
print("Error in message handler:", handler_exception)
else:
print(decoded_message)
except KeyboardInterrupt:
if self.verbose:
print("Received keyboard interrupt.")
self.close()
break
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)
break
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.")
def __enter__(self):
self._connect()
return self
def __exit__(self, exc_type, exc_value, traceback):
self.close()
yfinance/tickers.py
+12 -1
View File
@@ -22,6 +22,7 @@
from __future__ import print_function
from . import Ticker, multi
from .live import WebSocket
from .utils import print_once
from .data import YfData
from .const import _SENTINEL_
@@ -40,6 +41,9 @@ class Tickers:
self._data = YfData(session=session)
self._message_handler = None
self.ws = None
# self.tickers = _namedtuple(
# "Tickers", ticker_objects.keys(), rename=True
# )(*ticker_objects.values())
@@ -58,7 +62,7 @@ class Tickers:
return self.download(
period, interval,
start, end, prepost,
actions, auto_adjust, repair,
actions, auto_adjust, repair,
threads, group_by, progress,
timeout, **kwargs)
@@ -98,3 +102,10 @@ class Tickers:
def news(self):
return {ticker: [item for item in Ticker(ticker).news] for ticker in self.symbols}
def live(self, message_handler=None, verbose=True):
self._message_handler = message_handler
self.ws = WebSocket(verbose=verbose)
self.ws.subscribe(self.symbols)
self.ws.listen(self._message_handler)