From 332dad31384ecc7c6cbaca81c4e15d0660726ed2 Mon Sep 17 00:00:00 2001
From: moshferatu <gitea@steveharris.sh>
Date: Mon, 29 Jan 2024 07:18:05 -0800
Subject: [PATCH] Provide a nice wrapper to clients that streams options data
 into a pandas data frame

---
 client.py                       | 44 +++++++++++++++++++++++++++++++--
 get_options_chain_example.py    | 11 +++++++++
 option.py                       |  9 +++++++
 requirements.txt                |  1 +
 stream_options_chain_example.py |  6 ++++-
 5 files changed, 68 insertions(+), 3 deletions(-)
 create mode 100644 get_options_chain_example.py
 create mode 100644 option.py

diff --git a/client.py b/client.py
index 6740a15..5a26861 100644
--- a/client.py
+++ b/client.py
@@ -1,8 +1,11 @@
+import json
+import pandas as pd
 import requests
 import webbrowser
 
 from datetime import datetime
 from dotenv import load_dotenv
+from option import Option
 from os import getenv
 
 load_dotenv()
@@ -110,14 +113,51 @@ class TradeStationClient:
         response = requests.get(url, headers = headers, params = params, stream = True)
 
         try:
-            # TODO: Return a generator to the client.
             for line in response.iter_lines():
                 if line:
-                    print(line.decode('utf-8'))
+                    option_data = json.loads(line.decode('utf-8'))
+
+                    for leg in option_data.get('Legs', []):
+                        option = Option(
+                            Strike = leg.get('StrikePrice', 'StrikePrice'),
+                            Type = leg.get('OptionType', '').upper(),
+                            Bid = float(option_data.get('Bid', 0)),
+                            Ask = float(option_data.get('Ask', 0)),
+                            Delta = float(option_data.get('Delta', 0))
+                        )
+                        yield option
         except Exception as e:
             print(f'Error while streaming data: {str(e)}')
         finally:
             response.close()
+    
+    def get_options_chain(self, symbol: str, expiration: datetime = None, strike_proximity: int = 25) -> pd.DataFrame:
+        options_columns = ['Strike', 'Type', 'Bid', 'Ask', 'Delta']
+        dtypes = {'Strike': float, 'Type': str, 'Bid': float, 'Ask': float, 'Delta': float}
+        options_chain = pd.DataFrame(columns = options_columns).astype(dtypes)
+
+        for option in self.stream_options_chain(
+            symbol, expiration = expiration, strike_proximity = strike_proximity):
+            option_data = pd.DataFrame([{
+                'Strike': option.Strike,
+                'Type': option.Type,
+                'Bid': option.Bid,
+                'Ask': option.Ask,
+                'Delta': option.Delta
+            }])
+
+            # Only append a new row if the option doesn't already exist in the chain, otherwise update the row.
+            option_mask = (options_chain['Strike'] == option.Strike) & (options_chain['Type'] == option.Type)
+            if options_chain[option_mask].empty:
+                options_chain = pd.concat([options_chain, option_data], ignore_index = True)
+            else:
+                options_chain.loc[option_mask, ['Bid', 'Ask', 'Delta']] = [option.Bid, option.Ask, option.Delta]
+
+            # Number of strikes on either side of the spot price (x2) for each option type (x2).
+            if len(options_chain) >= strike_proximity * 4:
+                break
+
+        return options_chain
 
 if __name__ == '__main__':
     client = TradeStationClient()
diff --git a/get_options_chain_example.py b/get_options_chain_example.py
new file mode 100644
index 0000000..50c19c4
--- /dev/null
+++ b/get_options_chain_example.py
@@ -0,0 +1,11 @@
+from client import TradeStationClient
+from datetime import datetime
+from dotenv import load_dotenv
+from os import getenv
+
+load_dotenv()
+
+client = TradeStationClient(getenv('TRADESTATION_REFRESH_TOKEN'))
+
+options_chain = client.get_options_chain('$SPXW.X', expiration = datetime.now())
+print(options_chain)
\ No newline at end of file
diff --git a/option.py b/option.py
new file mode 100644
index 0000000..3502e3a
--- /dev/null
+++ b/option.py
@@ -0,0 +1,9 @@
+from dataclasses import dataclass
+
+@dataclass
+class Option:
+    Strike: str
+    Type: str
+    Bid: float
+    Ask: float
+    Delta: float
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 323ef8c..86a101e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,2 +1,3 @@
+pandas
 python-dotenv
 requests
\ No newline at end of file
diff --git a/stream_options_chain_example.py b/stream_options_chain_example.py
index 954ea55..60602a9 100644
--- a/stream_options_chain_example.py
+++ b/stream_options_chain_example.py
@@ -1,3 +1,5 @@
+import pandas as pd
+
 from client import TradeStationClient
 from datetime import datetime
 from dotenv import load_dotenv
@@ -6,4 +8,6 @@ from os import getenv
 load_dotenv()
 
 client = TradeStationClient(getenv('TRADESTATION_REFRESH_TOKEN'))
-client.stream_options_chain('$SPXW.X', expiration = datetime(2024, 1, 26))
\ No newline at end of file
+
+for option in client.stream_options_chain('$SPXW.X', expiration = datetime.now()):
+    print(option)
\ No newline at end of file