alpaca-py-nopandas 0.1.0__py3-none-any.whl

alpaca/data/live/websocket.py

@@ -0,0 +1,390 @@
+import asyncio
+import logging
+import queue
+from collections import defaultdict
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+import msgpack
+import websockets
+from pydantic import BaseModel
+from websockets.legacy import client as websockets_legacy
+
+from alpaca import __version__
+from alpaca.common.types import RawData
+from alpaca.data.models import (
+    Bar,
+    News,
+    Orderbook,
+    Quote,
+    Trade,
+    TradeCancel,
+    TradeCorrection,
+    TradingStatus,
+)
+
+log = logging.getLogger(__name__)
+
+
+class DataStream:
+    """
+    A base class for extracting out common functionality for data websockets
+    """
+
+    def __init__(
+        self,
+        endpoint: str,
+        api_key: str,
+        secret_key: str,
+        raw_data: bool = False,
+        websocket_params: Optional[Dict] = None,
+    ) -> None:
+        """Creates a new DataStream instance.
+
+        Args:
+            endpoint (str): The websocket endpoint to connect to
+            api_key (str): Alpaca API key.
+            secret_key (str): Alpaca API secret key.
+            raw_data (bool, optional): Whether to return raw API data or parsed data. Defaults to False.
+            websocket_params (Optional[Dict], optional): Any websocket connection configuration parameters. Defaults to None.
+        """
+        self._endpoint = endpoint
+        self._api_key = api_key
+        self._secret_key = secret_key
+        self._ws = None
+        self._running = False
+        self._loop = None
+        self._raw_data = raw_data
+        self._stop_stream_queue = queue.Queue()
+        self._handlers = {
+            "trades": {},
+            "quotes": {},
+            "orderbooks": {},
+            "bars": {},
+            "updatedBars": {},
+            "dailyBars": {},
+            "statuses": {},
+            "lulds": {},
+            "news": {},
+            "corrections": {},
+            "cancelErrors": {},
+        }
+        self._name = "data"
+        self._should_run = True
+        self._max_frame_size = 32768
+
+        self._websocket_params = {
+            "ping_interval": 10,
+            "ping_timeout": 180,
+            "max_queue": 1024,
+        }
+
+        if websocket_params:
+            self._websocket_params = websocket_params
+
+    async def _connect(self) -> None:
+        """Attempts to connect to the websocket endpoint.
+        If the connection attempt fails a value error is thrown.
+
+        Raises:
+            ValueError: Raised if there is an unsuccessful connection
+        """
+
+        extra_headers = {
+            "Content-Type": "application/msgpack",
+            "User-Agent": "APCA-PY/" + __version__,
+        }
+
+        log.info(f"connecting to {self._endpoint}")
+        self._ws = await websockets_legacy.connect(
+            self._endpoint,
+            extra_headers=extra_headers,
+            **self._websocket_params,
+        )
+        r = await self._ws.recv()
+        msg = msgpack.unpackb(r)
+        if msg[0]["T"] != "success" or msg[0]["msg"] != "connected":
+            raise ValueError("connected message not received")
+
+    async def _auth(self) -> None:
+        """Authenticates with API keys after a successful connection is established.
+
+        Raises:
+            ValueError: Raised if authentication is unsuccessful
+        """
+        await self._ws.send(
+            msgpack.packb(
+                {
+                    "action": "auth",
+                    "key": self._api_key,
+                    "secret": self._secret_key,
+                }
+            )
+        )
+        r = await self._ws.recv()
+        msg = msgpack.unpackb(r)
+        if msg[0]["T"] == "error":
+            raise ValueError(msg[0].get("msg", "auth failed"))
+        if msg[0]["T"] != "success" or msg[0]["msg"] != "authenticated":
+            raise ValueError("failed to authenticate")
+
+    async def _start_ws(self) -> None:
+        """Starts up the websocket connection. Attempts to connect to wss endpoint
+        and then authenticates with API keys.
+        """
+        await self._connect()
+        await self._auth()
+        log.info(f"connected to {self._endpoint}")
+
+    async def close(self) -> None:
+        """Closes the websocket connection."""
+        if self._ws:
+            await self._ws.close()
+            self._ws = None
+            self._running = False
+
+    async def stop_ws(self) -> None:
+        """Signals websocket connection should close by adding a closing message to the stop_stream_queue"""
+        self._should_run = False
+        if self._stop_stream_queue.empty():
+            self._stop_stream_queue.put_nowait({"should_stop": True})
+        await asyncio.sleep(0)
+
+    async def _consume(self) -> None:
+        """Distributes data from websocket connection to appropriate callbacks"""
+        while True:
+            if not self._stop_stream_queue.empty():
+                self._stop_stream_queue.get(timeout=1)
+                await self.close()
+                break
+            else:
+                try:
+                    r = await asyncio.wait_for(self._ws.recv(), 5)
+                    msgs = msgpack.unpackb(r)
+                    for msg in msgs:
+                        await self._dispatch(msg)
+                except asyncio.TimeoutError:
+                    # ws.recv is hanging when no data is received. by using
+                    # wait_for we break when no data is received, allowing us
+                    # to break the loop when needed
+                    pass
+
+    def _cast(self, msg: Dict) -> Union[BaseModel, RawData]:
+        """Parses data from websocket message if raw_data is False, otherwise
+        returns the raw websocket message.
+
+        Args:
+            msg (Dict): The message containing market data
+
+        Returns:
+            Union[BaseModel, RawData]: The raw or parsed message
+        """
+        if self._raw_data:
+            return msg
+        msg_type = msg.get("T")
+        if "t" in msg:
+            msg["t"] = msg["t"].to_datetime()
+        if msg_type == "n":
+            msg["created_at"] = msg["created_at"].to_datetime()
+            msg["updated_at"] = msg["updated_at"].to_datetime()
+            return News(msg)
+        if "S" not in msg:
+            return msg
+        if msg_type == "t":
+            return Trade(msg["S"], msg)
+        if msg_type == "q":
+            return Quote(msg["S"], msg)
+        if msg_type == "o":
+            return Orderbook(msg["S"], msg)
+        if msg_type in ("b", "u", "d"):
+            return Bar(msg["S"], msg)
+        if msg_type == "s":
+            return TradingStatus(msg["S"], msg)
+        if msg_type == "c":
+            return TradeCorrection(msg["S"], msg)
+        if msg_type == "x":
+            return TradeCancel(msg["S"], msg)
+        return msg
+
+    async def _dispatch(self, msg: Dict) -> None:
+        """Distributes the message from websocket connection to the appropriate handler.
+
+        Args:
+            msg (Dict): The message from the websocket connection
+        """
+        msg_type = msg.get("T")
+        if msg_type == "subscription":
+            sub = [f"{k}: {msg.get(k, [])}" for k in self._handlers if msg.get(k)]
+            log.info(f'subscribed to {", ".join(sub)}')
+            return
+
+        if msg_type == "error":
+            log.error(f'error: {msg.get("msg")} ({msg.get("code")})')
+            return
+
+        if msg_type == "n":
+            symbols = msg.get("symbols", "*")
+            star_handler_called = False
+            handlers_to_call = []
+            news = self._cast(msg)
+            for symbol in set(symbols):
+                if symbol in self._handlers["news"]:
+                    handler = self._handlers["news"].get(symbol)
+                elif not star_handler_called:
+                    handler = self._handlers["news"].get("*")
+                    star_handler_called = True
+                else:
+                    handler = None
+                if handler:
+                    handlers_to_call.append(handler(news))
+            if handlers_to_call:
+                await asyncio.gather(*handlers_to_call)
+            return
+
+        channel_types = {
+            "t": "trades",
+            "q": "quotes",
+            "o": "orderbooks",
+            "b": "bars",
+            "u": "updatedBars",
+            "d": "dailyBars",
+            "s": "statuses",
+            "l": "lulds",
+            "n": "news",
+            "c": "corrections",
+            "x": "cancelErrors",
+        }
+        channel = channel_types.get(msg_type)
+        if not channel:
+            return
+        symbol = msg.get("S")
+        handler = self._handlers[channel].get(symbol, self._handlers[channel].get("*"))
+        if handler:
+            await handler(self._cast(msg))
+
+    def _subscribe(
+        self, handler: Callable, symbols: Tuple[str], handlers: Dict
+    ) -> None:
+        """Subscribes a coroutine callback function to receive data for a tuple of symbols
+
+        Args:
+            handler (Callable): The coroutine callback function to receive data
+            symbols (Tuple[str]): The tuple containing the symbols to be subscribed to
+            handlers (Dict): The dictionary of coroutine callback functions keyed by symbol
+        """
+        self._ensure_coroutine(handler)
+        for symbol in symbols:
+            handlers[symbol] = handler
+        if self._running:
+            asyncio.run_coroutine_threadsafe(
+                self._send_subscribe_msg(), self._loop
+            ).result()
+
+    async def _send_subscribe_msg(self) -> None:
+        msg = defaultdict(list)
+        for k, v in self._handlers.items():
+            if k not in ("cancelErrors", "corrections") and v:
+                for s in v.keys():
+                    msg[k].append(s)
+        msg["action"] = "subscribe"
+        bs = msgpack.packb(msg)
+        frames = (
+            bs[i : i + self._max_frame_size]
+            for i in range(0, len(bs), self._max_frame_size)
+        )
+        await self._ws.send(frames)
+
+    def _unsubscribe(self, channel: str, symbols: List[str]) -> None:
+        if self._running:
+            asyncio.run_coroutine_threadsafe(
+                self._send_unsubscribe_msg(channel, symbols), self._loop
+            ).result()
+        for symbol in symbols:
+            del self._handlers[channel][symbol]
+
+    async def _send_unsubscribe_msg(self, channel: str, symbols: List[str]) -> None:
+        if symbols:
+            await self._ws.send(
+                msgpack.packb(
+                    {
+                        "action": "unsubscribe",
+                        channel: symbols,
+                    }
+                )
+            )
+
+    async def _run_forever(self) -> None:
+        """Starts event loop for receiving data from websocket connection and handles
+        distributing messages
+        """
+        self._loop = asyncio.get_running_loop()
+        # do not start the websocket connection until we subscribe to something
+        while not any(
+            v
+            for k, v in self._handlers.items()
+            if k not in ("cancelErrors", "corrections")
+        ):
+            if not self._stop_stream_queue.empty():
+                # the ws was signaled to stop before starting the loop so
+                # we break
+                self._stop_stream_queue.get(timeout=1)
+                return
+            await asyncio.sleep(0)
+        log.info(f"started {self._name} stream")
+        self._should_run = True
+        self._running = False
+        while True:
+            try:
+                if not self._should_run:
+                    # when signaling to stop, this is how we break run_forever
+                    log.info("{} stream stopped".format(self._name))
+                    return
+                if not self._running:
+                    log.info("starting {} websocket connection".format(self._name))
+                    await self._start_ws()
+                    await self._send_subscribe_msg()
+                    self._running = True
+                await self._consume()
+            except websockets.WebSocketException as wse:
+                await self.close()
+                self._running = False
+                log.warning("data websocket error, restarting connection: " + str(wse))
+            except ValueError as ve:
+                if "insufficient subscription" in str(ve):
+                    await self.close()
+                    self._running = False
+                    log.exception(f"error during websocket communication: {str(ve)}")
+                    return
+                log.exception(f"error during websocket communication: {str(ve)}")
+            except Exception as e:
+                log.exception(f"error during websocket communication: {str(e)}")
+            finally:
+                await asyncio.sleep(0)
+
+    def run(self) -> None:
+        """Starts up the websocket connection's event loop"""
+        try:
+            asyncio.run(self._run_forever())
+        except KeyboardInterrupt:
+            print("keyboard interrupt, bye")
+            pass
+        finally:
+            self.stop()
+
+    def stop(self) -> None:
+        """Stops the websocket connection."""
+        if self._loop.is_running():
+            asyncio.run_coroutine_threadsafe(self.stop_ws(), self._loop).result(
+                timeout=5
+            )
+
+    def _ensure_coroutine(self, handler: Callable) -> None:
+        """Checks if a method is an asyncio coroutine method
+
+        Args:
+            handler (Callable): A method to be checked for coroutineness
+
+        Raises:
+            ValueError: Raised if the input method is not a coroutine
+        """
+        if not asyncio.iscoroutinefunction(handler):
+            raise ValueError("handler must be a coroutine function")

alpaca/data/mappings.py

@@ -0,0 +1,84 @@
+from typing import Dict
+
+BAR_MAPPING: Dict[str, str] = {
+    "t": "timestamp",
+    "o": "open",
+    "h": "high",
+    "l": "low",
+    "c": "close",
+    "v": "volume",
+    "n": "trade_count",
+    "vw": "vwap",
+}
+
+QUOTE_MAPPING: Dict[str, str] = {
+    "t": "timestamp",
+    "ax": "ask_exchange",
+    "ap": "ask_price",
+    "as": "ask_size",
+    "bx": "bid_exchange",
+    "bp": "bid_price",
+    "bs": "bid_size",
+    "c": "conditions",
+    "z": "tape",
+}
+
+TRADE_MAPPING: Dict[str, str] = {
+    "t": "timestamp",
+    "p": "price",
+    "s": "size",
+    "x": "exchange",
+    "i": "id",
+    "c": "conditions",
+    "z": "tape",
+}
+
+SNAPSHOT_MAPPING: Dict[str, str] = {
+    "latestTrade": "latest_trade",
+    "latestQuote": "latest_quote",
+    "minuteBar": "minute_bar",
+    "dailyBar": "daily_bar",
+    "prevDailyBar": "previous_daily_bar",
+    "impliedVolatility": "implied_volatility",
+    "greeks": "greeks",
+}
+
+ORDERBOOK_MAPPING: Dict[str, str] = {
+    "t": "timestamp",
+    "b": "bids",
+    "a": "asks",
+    "r": "reset",
+}
+
+TRADING_STATUS_MAPPING: Dict[str, str] = {
+    "t": "timestamp",
+    "sc": "status_code",
+    "sm": "status_message",
+    "rc": "reason_code",
+    "rm": "reason_message",
+    "z": "tape",
+}
+
+TRADE_CANCEL_MAPPING: Dict[str, str] = {
+    "t": "timestamp",
+    "p": "price",
+    "s": "size",
+    "x": "exchange",
+    "i": "id",
+    "a": "action",
+    "z": "tape",
+}
+
+TRADE_CORRECTION_MAPPING: Dict[str, str] = {
+    "t": "timestamp",
+    "x": "exchange",
+    "oi": "original_id",
+    "op": "original_price",
+    "os": "original_size",
+    "oc": "original_conditions",
+    "ci": "corrected_id",
+    "cp": "corrected_price",
+    "cs": "corrected_size",
+    "cc": "corrected_conditions",
+    "z": "tape",
+}

alpaca/data/models/__init__.py

@@ -0,0 +1,7 @@
+from alpaca.data.models.bars import *
+from alpaca.data.models.orderbooks import *
+from alpaca.data.models.news import *
+from alpaca.data.models.quotes import *
+from alpaca.data.models.trades import *
+from alpaca.data.models.snapshots import *
+from alpaca.data.models.orderbooks import *

alpaca/data/models/bars.py

@@ -0,0 +1,83 @@
+from datetime import datetime
+from typing import Dict, List, Optional
+
+from alpaca.common.models import ValidateBaseModel as BaseModel
+from alpaca.common.types import RawData
+from alpaca.data.mappings import BAR_MAPPING
+from alpaca.data.models.base import BaseDataSet, TimeSeriesMixin
+
+
+class Bar(BaseModel):
+    """Represents one bar/candlestick of aggregated trade data over a specified interval.
+
+    Attributes:
+        symbol (str): The ticker identifier for the security whose data forms the bar.
+        timestamp (datetime): The opening timestamp of the bar.
+        open (float): The opening price of the interval.
+        high (float): The high price during the interval.
+        low (float): The low price during the interval.
+        close (float): The closing price of the interval.
+        volume (float): The volume traded over the interval.
+        trade_count (Optional[float]): The number of trades that occurred.
+        vwap (Optional[float]): The volume weighted average price.
+        exchange (Optional[float]): The exchange the bar was formed on.
+    """
+
+    symbol: str
+    timestamp: datetime
+    open: float
+    high: float
+    low: float
+    close: float
+    volume: float
+    trade_count: Optional[float]
+    vwap: Optional[float]
+
+    def __init__(self, symbol: str, raw_data: RawData) -> None:
+        """Instantiates a bar
+
+        Args:
+            raw_data (RawData): Raw unparsed bar data from API, contains ohlc and other fields.
+        """
+        mapped_bar = {}
+
+        if raw_data is not None:
+            mapped_bar = {
+                BAR_MAPPING[key]: val
+                for key, val in raw_data.items()
+                if key in BAR_MAPPING
+            }
+
+        super().__init__(symbol=symbol, **mapped_bar)
+
+
+class BarSet(BaseDataSet, TimeSeriesMixin):
+    """A collection of Bars.
+
+    Attributes:
+        data (Dict[str, List[Bar]]): The collection of Bars keyed by symbol.
+    """
+
+    data: Dict[str, List[Bar]] = {}
+
+    def __init__(
+        self,
+        raw_data: RawData,
+    ) -> None:
+        """A collection of Bars.
+
+        Args:
+            raw_data (RawData): The collection of raw bar data from API keyed by Symbol.
+        """
+
+        parsed_bars = {}
+
+        raw_bars = raw_data
+
+        if raw_bars is not None:
+            for symbol, bars in raw_bars.items():
+                parsed_bars[symbol] = [
+                    Bar(symbol, bar) for bar in bars if bar is not None
+                ]
+
+        super().__init__(data=parsed_bars)

alpaca/data/models/base.py

@@ -0,0 +1,45 @@
+from typing import Any, Dict, List
+
+from alpaca.common.models import ValidateBaseModel as BaseModel
+
+
+class TimeSeriesMixin:
+    pass
+
+
+class BaseDataSet(BaseModel):
+    """
+    Base class to process data models for trades, bars quotes, and news.
+    """
+
+    data: Dict[str, List[BaseModel]] = {}
+
+    def __getitem__(self, symbol: str) -> Any:
+        """Gives dictionary-like access to multi-symbol data
+
+        Args:
+            symbol (str): The ticker identifier for the desired data
+
+        Raises:
+            KeyError: Symbol does not exist for data
+
+        Returns:
+            List[Bar]: The data for the given symbol
+        """
+        if symbol not in self.data:
+            raise KeyError(f"No key {symbol} was found.")
+
+        return self.data[symbol]
+
+    def dict(self, **kwargs) -> dict:
+        """
+        Gives dictionary representation of data.
+
+        Returns:
+            dict: The data in dictionary form.
+        """
+        # converts each data (Bar, Quote, etc) in the symbol specific lists to its dictionary format
+        return {
+            symbol: list(map(lambda d: d.model_dump(), data_list))
+            for symbol, data_list in self.data.items()
+        }