From cee81dd3a422b03f09e299c14cc73e65b9a551da Mon Sep 17 00:00:00 2001
From: moshferatu <gitea@steveharris.sh>
Date: Mon, 9 Oct 2023 13:09:02 -0700
Subject: [PATCH] Add support for streaming market data (quotes)

---
 .gitignore                        |  4 +-
 requirements.txt                  |  4 +-
 streaming_data_example.py         | 12 ++++++
 streaming_option_chain_example.py | 27 ++++++++++++
 tastytrade/tastytrade.py          | 69 ++++++++++++++++++++++++++++++-
 5 files changed, 113 insertions(+), 3 deletions(-)
 create mode 100644 streaming_data_example.py
 create mode 100644 streaming_option_chain_example.py

diff --git a/.gitignore b/.gitignore
index 2eea525..5d84dee 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,3 @@
-.env
\ No newline at end of file
+.env
+*.egg-info/
+__pycache__/
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 663bd1f..f165439 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1 +1,3 @@
-requests
\ No newline at end of file
+dotenv
+requests
+websockets
\ No newline at end of file
diff --git a/streaming_data_example.py b/streaming_data_example.py
new file mode 100644
index 0000000..a76af19
--- /dev/null
+++ b/streaming_data_example.py
@@ -0,0 +1,12 @@
+from dotenv import load_dotenv
+from os import getenv
+from tastytrade import Tastytrade
+
+load_dotenv()
+account = getenv("TASTYTRADE_ACCOUNT")
+username = getenv("TASTYTRADE_USERNAME")
+password = getenv("TASTYTRADE_PASSWORD")
+
+client = Tastytrade(username, password)
+client.login()
+client.start_streaming(['AAPL', 'AMZN'])
\ No newline at end of file
diff --git a/streaming_option_chain_example.py b/streaming_option_chain_example.py
new file mode 100644
index 0000000..e0e23a5
--- /dev/null
+++ b/streaming_option_chain_example.py
@@ -0,0 +1,27 @@
+"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
+    This example does not currently work.
+    The DXLink connection does not currently support options data.
+    Refer to the comments on this video: https://www.youtube.com/watch?v=qHL4Jy6yIC8
+"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
+
+from datetime import datetime
+from dotenv import load_dotenv
+from os import getenv
+from tastytrade import Tastytrade
+
+load_dotenv()
+account = getenv("TASTYTRADE_ACCOUNT")
+username = getenv("TASTYTRADE_USERNAME")
+password = getenv("TASTYTRADE_PASSWORD")
+
+client = Tastytrade(username, password)
+client.login()
+
+current_date = datetime.now().strftime('%y%m%d')
+symbol_prefix = f'.SPXW{current_date}'
+
+option_chain_data = client.get_option_chain_compact('SPX')['data']['items'][0]
+streamer_symbols = option_chain_data['streamer-symbols']
+zero_dte_symbols = [symbol for symbol in streamer_symbols if symbol.startswith(symbol_prefix)]
+
+client.start_streaming(zero_dte_symbols)
\ No newline at end of file
diff --git a/tastytrade/tastytrade.py b/tastytrade/tastytrade.py
index 5dd8d57..eea2581 100644
--- a/tastytrade/tastytrade.py
+++ b/tastytrade/tastytrade.py
@@ -1,5 +1,7 @@
+import asyncio
 import json
 import requests
+import websockets
 
 class Tastytrade:
 
@@ -47,4 +49,69 @@ class Tastytrade:
         return self.get(path = f'/accounts/{account_number}/positions').json()
     
     def submit_order(self, account_number: str, order: dict) -> dict:
-        return self.post(path = f'/accounts/{account_number}/orders', data = order).json()
\ No newline at end of file
+        return self.post(path = f'/accounts/{account_number}/orders', data = order).json()
+    
+    async def setup_connection(self, ws):
+        setup_message = {
+            'type': 'SETUP',
+            'channel': 0,
+            'keepaliveTimeout': 60,
+            'acceptKeepaliveTimeout': 60,
+            'version': '0.1-js/1.0.0'
+        }
+        await ws.send(json.dumps(setup_message))
+        response = await ws.recv()
+        return json.loads(response)
+
+    async def authenticate(self, ws, token):
+        auth_message = {
+            'type': 'AUTH',
+            'channel': 0,
+            'token': token
+        }
+        await ws.send(json.dumps(auth_message))
+        response = await ws.recv()
+        return json.loads(response)
+
+    async def stream_market_data(self, symbols: list) -> None:
+        token_response = self.get('/api-quote-tokens').json()
+        api_quote_token = token_response['data']['token']
+        dxlink_url = token_response['data']['dxlink-url']
+
+        async with websockets.connect(dxlink_url) as ws:
+            # Documentation: https://demo.dxfeed.com/dxlink-ws/debug/#/protocol
+            setup_response = await self.setup_connection(ws)
+            auth_response = await self.authenticate(ws, api_quote_token)
+
+            # For some reason, the auth response is always unauthorized yet subsequent requests work.
+            # if auth_response.get('state') != 'AUTHORIZED':
+            #     print('Authorization failed.')
+            #     return
+
+            # Request new channel for Quote (and Greeks) events
+            channel_request = {
+                'type': 'CHANNEL_REQUEST',
+                'channel': 1,
+                'service': 'FEED',
+                'parameters': {
+                    'contract': 'AUTO'
+                }
+            }
+            await ws.send(json.dumps(channel_request))
+            channel_response = await ws.recv()
+
+            subscription_message = {
+                'type': 'FEED_SUBSCRIPTION',
+                'channel': 1,
+                'add': [{'symbol': symbol, 'type': 'Quote'} for symbol in symbols]
+            }
+            await ws.send(json.dumps(subscription_message))
+            
+            async for message in ws:
+                data = json.loads(message)
+                if data["type"] == "FEED_DATA":
+                    # TODO: Have the caller provide a callback function to handle the data.
+                    print(data)
+
+    def start_streaming(self, symbols: list):
+        asyncio.get_event_loop().run_until_complete(self.stream_market_data(symbols))
\ No newline at end of file