tradestation/client.py

103 lines
3.4 KiB
Python

import requests
import webbrowser
from dotenv import load_dotenv
from os import getenv
load_dotenv()
# TODO: Add support for paper trading API endpoint.
class TradeStationClient:
def __init__(self, refresh_token: str = None) -> None:
self.id = getenv('TRADESTATION_CLIENT_ID')
self.redirect_uri = getenv('TRADESTATION_REDIRECT_URI')
self.scope = getenv('TRADESTATION_CLIENT_SCOPE')
self.secret = getenv('TRADESTATION_CLIENT_SECRET')
self.state = getenv('TRADESTATION_CLIENT_STATE')
# Must be retrieved via authorization flow.
# TODO: Automate authorization.
self.access_token = None
self.id_token = None
self.refresh_token = None
# For bypassing normal authorization flow.
if refresh_token:
self.refresh_token = refresh_token
self.refresh_access_token()
def open_authorization_url(self):
"""
Open the TradeStation authorization URL in the default web browser.
"""
url = (
f"https://signin.tradestation.com/authorize?"
f"response_type=code&"
f"client_id={self.id}&"
f"redirect_uri={self.redirect_uri}&"
f"audience=https://api.tradestation.com&"
f"state={self.state}&"
f"scope={self.scope}"
)
webbrowser.open(url)
def get_tokens(self, authorization_code):
"""
Exchange the authorization code for access token, ID token, and refresh token.
"""
url = 'https://signin.tradestation.com/oauth/token'
headers = {'content-type': 'application/x-www-form-urlencoded'}
payload = {
'grant_type': 'authorization_code',
'client_id': self.id,
'client_secret': self.secret,
'code': authorization_code,
'redirect_uri': self.redirect_uri
}
response = requests.post(url, headers=headers, data=payload)
if response.status_code == 200:
return response.json()
else:
raise Exception(f"Failed to get tokens: {response.text}")
def refresh_access_token(self):
"""
Use the refresh token to obtain a new access token.
"""
if not self.refresh_token:
raise Exception("Refresh token is not available.")
url = 'https://signin.tradestation.com/oauth/token'
headers = {'content-type': 'application/x-www-form-urlencoded'}
payload = {
'grant_type': 'refresh_token',
'client_id': self.id,
'refresh_token': self.refresh_token,
'client_secret': self.secret
}
response = requests.post(url, headers=headers, data=payload)
if response.status_code == 200:
tokens = response.json()
self.access_token = tokens.get('access_token')
self.id_token = tokens.get('id_token')
else:
raise Exception(f"Failed to refresh token: {response.text}")
if __name__ == '__main__':
client = TradeStationClient()
client.open_authorization_url()
authorization_code = input("Please enter the authorization code you received: ")
try:
tokens = client.get_tokens(authorization_code)
print("Access Token:", tokens['access_token'])
print("ID Token:", tokens['id_token'])
print("Refresh Token:", tokens['refresh_token'])
except Exception as e:
print(str(e))