From 3180feef1963c2616aa3f5a2fff6002462b98f91 Mon Sep 17 00:00:00 2001
From: Graeme22 <graeme.holliday@pm.me>
Date: Tue, 30 May 2023 13:53:25 -0500
Subject: [PATCH] alert streamer mapping, small fixes

---
 README.rst                       |  2 +-
 tastytrade/__init__.py           |  2 +-
 tastytrade/dxfeed/timeandsale.py |  6 +++-
 tastytrade/streamer.py           | 49 ++++++++++++++++++++++++++++----
 4 files changed, 51 insertions(+), 8 deletions(-)

diff --git a/README.rst b/README.rst
index b89e2fb..a13992a 100644
--- a/README.rst
+++ b/README.rst
@@ -15,7 +15,7 @@ Tastytrade Python SDK
 
 .. inclusion-marker
 
-A simple, reverse-engineered SDK for Tastytrade built on their (mostly) public API. This will allow you to create trading algorithms for whatever strategies you may have quickly and painlessly in Python.
+A simple, reverse-engineered SDK for Tastytrade built on their (now mostly public) API. This will allow you to create trading algorithms for whatever strategies you may have quickly and painlessly in Python.
 
 Installation
 ------------
diff --git a/tastytrade/__init__.py b/tastytrade/__init__.py
index 8db87f2..3fadb64 100644
--- a/tastytrade/__init__.py
+++ b/tastytrade/__init__.py
@@ -2,7 +2,7 @@
 
 API_URL = 'https://api.tastyworks.com'
 CERT_URL = 'https://api.cert.tastyworks.com'
-VERSION = '5.2'
+VERSION = '5.3'
 
 
 logger = logging.getLogger(__name__)
diff --git a/tastytrade/dxfeed/timeandsale.py b/tastytrade/dxfeed/timeandsale.py
index 4d22c42..6bd0e5a 100644
--- a/tastytrade/dxfeed/timeandsale.py
+++ b/tastytrade/dxfeed/timeandsale.py
@@ -6,7 +6,11 @@
 @dataclass
 class TimeAndSale(Event):
     """
-    TimeAndSale event represents a trade or other market event with a price, like market open/close price. TimeAndSale events are intended to provide information about trades in a continuous-time slice (unlike Trade events which are supposed to provide snapshots about the most recent trade). TimeAndSale events have a unique index that can be used for later correction/cancellation processing.
+    TimeAndSale event represents a trade or other market event with a price, like
+    market open/close price. TimeAndSale events are intended to provide information
+    about trades in a continuous-time slice (unlike Trade events which are supposed
+    to provide snapshots about the most recent trade). TimeAndSale events have a
+    unique index that can be used for later correction/cancellation processing.
     """
     #: symbol of this event
     eventSymbol: str
diff --git a/tastytrade/streamer.py b/tastytrade/streamer.py
index 598b695..694a76a 100644
--- a/tastytrade/streamer.py
+++ b/tastytrade/streamer.py
@@ -9,7 +9,8 @@
 import websockets
 
 from tastytrade import logger
-from tastytrade.account import Account
+from tastytrade.account import (Account, AccountBalance, CurrentPosition,
+                                TradingStatus)
 from tastytrade.dxfeed import Channel
 from tastytrade.dxfeed.candle import Candle
 from tastytrade.dxfeed.event import Event, EventType
@@ -20,6 +21,7 @@
 from tastytrade.dxfeed.theoprice import TheoPrice
 from tastytrade.dxfeed.timeandsale import TimeAndSale
 from tastytrade.dxfeed.trade import Trade
+from tastytrade.order import PlacedOrder
 from tastytrade.session import Session
 from tastytrade.utils import TastytradeError, validate_response
 
@@ -102,14 +104,51 @@ async def _connect(self) -> None:
                 logger.debug('raw message: %s', raw_message)
                 await self._queue.put(json.loads(raw_message))
 
-    async def listen(self) -> AsyncIterator[Any]:
+    async def listen(self) -> AsyncIterator[Union[
+        AccountBalance,
+        CurrentPosition,
+        PlacedOrder,
+        TradingStatus,
+        dict  # some possible messages are not yet implemented
+    ]]:
         """
-        Iterate over non-heartbeat messages received from the streamer.
+        Iterate over non-heartbeat messages received from the streamer,
+        mapping them to their appropriate data class.
         """
         while True:
             data = await self._queue.get()
-            if data.get('action') != SubscriptionType.HEARTBEAT:
-                yield data
+            type_str = data.get('type')
+            if type_str is not None:
+                yield self._map_message(type_str, data['data'])
+            elif data.get('action') != 'heartbeat':
+                logger.debug('subscription message: %s', data)
+
+    def _map_message(self, type_str: str, data: dict) -> Union[
+        AccountBalance,
+        CurrentPosition,
+        PlacedOrder,
+        TradingStatus,
+        dict  # some possible messages are not yet implemented
+    ]:
+        """
+        TODO: implement the following:
+        - OrderChain
+        - UnderlyingYearGainSummary
+        - User status related messages
+        - Watchlist related messages
+        - Quote alert messages
+        - Others?
+        """
+        if type_str == 'AccountBalance':
+            return AccountBalance(**data)
+        elif type_str == 'CurrentPosition':
+            return CurrentPosition(**data)
+        elif type_str == 'Order':
+            return PlacedOrder(**data)
+        elif type_str == 'TradingStatus':
+            return TradingStatus(**data)
+        else:
+            return data
 
     async def account_subscribe(self, accounts: list[Account]) -> None:
         """