Mercurial

comparison schwab_trader/schwab_client.py @ 220:eb8b4230fdb9

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
[schwab-trader] Add guarded trading experiment
author MrJuneJune <me@mrjunejune.com>
date Sun, 02 Aug 2026 08:52:13 -0700
parents
children
comparison
equal deleted inserted replaced
214:4c725fde6999 220:eb8b4230fdb9
1 from __future__ import annotations
2
3 import base64
4 import json
5 import os
6 import stat
7 import time
8 import urllib.error
9 import urllib.parse
10 import urllib.request
11 from dataclasses import dataclass
12 from pathlib import Path
13 from typing import Any
14
15
16 AUTH_URL = "https://api.schwabapi.com/v1/oauth/authorize"
17 TOKEN_URL = "https://api.schwabapi.com/v1/oauth/token"
18 TRADER_BASE_URL = "https://api.schwabapi.com/trader/v1"
19 DEFAULT_TOKEN_FILE = "~/.config/zenbu/schwab_tokens.json"
20
21
22 class SchwabError(RuntimeError):
23 pass
24
25
26 @dataclass(frozen=True)
27 class SchwabConfig:
28 app_key: str
29 app_secret: str
30 redirect_uri: str
31 token_file: Path
32
33 @classmethod
34 def from_env(cls) -> "SchwabConfig":
35 app_key = os.environ.get("SCHWAB_APP_KEY", "").strip()
36 app_secret = os.environ.get("SCHWAB_APP_SECRET", "").strip()
37 redirect_uri = os.environ.get("SCHWAB_REDIRECT_URI", "").strip()
38 token_file = Path(os.environ.get("SCHWAB_TOKEN_FILE", DEFAULT_TOKEN_FILE)).expanduser()
39
40 missing = [
41 name
42 for name, value in (
43 ("SCHWAB_APP_KEY", app_key),
44 ("SCHWAB_APP_SECRET", app_secret),
45 ("SCHWAB_REDIRECT_URI", redirect_uri),
46 )
47 if not value
48 ]
49 if missing:
50 raise SchwabError("Missing required environment variables: " + ", ".join(missing))
51
52 return cls(
53 app_key=app_key,
54 app_secret=app_secret,
55 redirect_uri=redirect_uri,
56 token_file=token_file,
57 )
58
59
60 @dataclass(frozen=True)
61 class ApiResponse:
62 status: int
63 headers: dict[str, str]
64 body: Any
65 raw_body: str
66
67
68 def build_authorization_url(app_key: str, redirect_uri: str, state: str | None = None) -> str:
69 params = {
70 "response_type": "code",
71 "client_id": app_key,
72 "redirect_uri": redirect_uri,
73 }
74 if state:
75 params["state"] = state
76 return AUTH_URL + "?" + urllib.parse.urlencode(params)
77
78
79 def extract_authorization_code(code_or_url: str) -> str:
80 value = code_or_url.strip()
81 if not value:
82 raise SchwabError("Authorization code is empty")
83
84 parsed = urllib.parse.urlparse(value)
85 if parsed.scheme and parsed.netloc:
86 query = urllib.parse.parse_qs(parsed.query)
87 codes = query.get("code")
88 if not codes or not codes[0]:
89 raise SchwabError("No code= parameter found in callback URL")
90 return codes[0]
91
92 return urllib.parse.unquote(value)
93
94
95 def exchange_code_for_tokens(config: SchwabConfig, code_or_url: str) -> dict[str, Any]:
96 code = extract_authorization_code(code_or_url)
97 return _token_request(
98 config,
99 {
100 "grant_type": "authorization_code",
101 "code": code,
102 "redirect_uri": config.redirect_uri,
103 },
104 )
105
106
107 def refresh_tokens(config: SchwabConfig, refresh_token: str | None = None) -> dict[str, Any]:
108 token_value = refresh_token
109 if token_value is None:
110 existing = load_tokens(config.token_file)
111 token_value = existing.get("refresh_token")
112 if not token_value:
113 raise SchwabError("No refresh token available")
114
115 return _token_request(
116 config,
117 {
118 "grant_type": "refresh_token",
119 "refresh_token": token_value,
120 },
121 )
122
123
124 def save_tokens(path: Path, tokens: dict[str, Any]) -> None:
125 payload = dict(tokens)
126 payload["saved_at"] = int(time.time())
127 path.parent.mkdir(parents=True, exist_ok=True)
128 path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
129 path.chmod(stat.S_IRUSR | stat.S_IWUSR)
130
131
132 def load_tokens(path: Path) -> dict[str, Any]:
133 if not path.exists():
134 raise SchwabError(f"Token file does not exist: {path}")
135 return json.loads(path.read_text(encoding="utf-8"))
136
137
138 def get_account_numbers(access_token: str) -> ApiResponse:
139 return _api_request("GET", "/accounts/accountNumbers", access_token)
140
141
142 def get_accounts(access_token: str, fields: str | None = None) -> ApiResponse:
143 query = ""
144 if fields:
145 query = "?" + urllib.parse.urlencode({"fields": fields})
146 return _api_request("GET", "/accounts" + query, access_token)
147
148
149 def get_account(access_token: str, account_hash: str, fields: str | None = None) -> ApiResponse:
150 query = ""
151 if fields:
152 query = "?" + urllib.parse.urlencode({"fields": fields})
153 return _api_request("GET", f"/accounts/{urllib.parse.quote(account_hash)}{query}", access_token)
154
155
156 def build_equity_order(
157 action: str,
158 symbol: str,
159 quantity: float,
160 order_type: str = "MARKET",
161 price: float | None = None,
162 duration: str = "DAY",
163 session: str = "NORMAL",
164 ) -> dict[str, Any]:
165 normalized_action = action.upper()
166 normalized_symbol = symbol.upper()
167 normalized_order_type = order_type.upper()
168 normalized_duration = duration.upper()
169 normalized_session = session.upper()
170
171 if normalized_action not in {"BUY", "SELL"}:
172 raise SchwabError("action must be BUY or SELL")
173 if not normalized_symbol:
174 raise SchwabError("symbol is required")
175 if quantity <= 0:
176 raise SchwabError("quantity must be greater than zero")
177 if normalized_order_type not in {"MARKET", "LIMIT"}:
178 raise SchwabError("order_type must be MARKET or LIMIT")
179 if normalized_order_type == "LIMIT" and price is None:
180 raise SchwabError("LIMIT orders require --price")
181 if normalized_order_type == "MARKET" and price is not None:
182 raise SchwabError("MARKET orders cannot include --price")
183
184 order: dict[str, Any] = {
185 "orderType": normalized_order_type,
186 "session": normalized_session,
187 "duration": normalized_duration,
188 "orderStrategyType": "SINGLE",
189 "orderLegCollection": [
190 {
191 "instruction": normalized_action,
192 "quantity": quantity,
193 "instrument": {
194 "symbol": normalized_symbol,
195 "assetType": "EQUITY",
196 },
197 }
198 ],
199 }
200 if price is not None:
201 order["price"] = f"{price:.2f}"
202 return order
203
204
205 def place_order(access_token: str, account_hash: str, order: dict[str, Any]) -> ApiResponse:
206 path = f"/accounts/{urllib.parse.quote(account_hash)}/orders"
207 return _api_request("POST", path, access_token, order)
208
209
210 def _token_request(config: SchwabConfig, form: dict[str, str]) -> dict[str, Any]:
211 credentials = f"{config.app_key}:{config.app_secret}".encode("utf-8")
212 headers = {
213 "Authorization": "Basic " + base64.b64encode(credentials).decode("ascii"),
214 "Content-Type": "application/x-www-form-urlencoded",
215 "Accept": "application/json",
216 }
217 data = urllib.parse.urlencode(form).encode("utf-8")
218 request = urllib.request.Request(TOKEN_URL, data=data, headers=headers, method="POST")
219 response = _open_request(request)
220 if not isinstance(response.body, dict):
221 raise SchwabError("Token endpoint did not return a JSON object")
222 return response.body
223
224
225 def _api_request(
226 method: str,
227 path: str,
228 access_token: str,
229 body: dict[str, Any] | None = None,
230 ) -> ApiResponse:
231 headers = {
232 "Authorization": f"Bearer {access_token}",
233 "Accept": "application/json",
234 }
235 data = None
236 if body is not None:
237 data = json.dumps(body).encode("utf-8")
238 headers["Content-Type"] = "application/json"
239
240 request = urllib.request.Request(
241 TRADER_BASE_URL + path,
242 data=data,
243 headers=headers,
244 method=method,
245 )
246 return _open_request(request)
247
248
249 def _open_request(request: urllib.request.Request) -> ApiResponse:
250 try:
251 with urllib.request.urlopen(request, timeout=30) as response:
252 raw = response.read().decode("utf-8")
253 return ApiResponse(
254 status=response.status,
255 headers=dict(response.headers.items()),
256 body=_parse_json_or_text(raw),
257 raw_body=raw,
258 )
259 except urllib.error.HTTPError as error:
260 raw = error.read().decode("utf-8", errors="replace")
261 raise SchwabError(
262 f"Schwab API request failed with HTTP {error.code}: {raw or error.reason}"
263 ) from error
264 except urllib.error.URLError as error:
265 raise SchwabError(f"Schwab API request failed: {error.reason}") from error
266
267
268 def _parse_json_or_text(raw: str) -> Any:
269 if not raw:
270 return None
271 try:
272 return json.loads(raw)
273 except json.JSONDecodeError:
274 return raw
275