diff schwab_trader/schwab_client.py @ 220:eb8b4230fdb9

[schwab-trader] Add guarded trading experiment
author MrJuneJune <me@mrjunejune.com>
date Sun, 02 Aug 2026 08:52:13 -0700
parents
children
line wrap: on
line diff
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/schwab_trader/schwab_client.py	Sun Aug 02 08:52:13 2026 -0700
@@ -0,0 +1,275 @@
+from __future__ import annotations
+
+import base64
+import json
+import os
+import stat
+import time
+import urllib.error
+import urllib.parse
+import urllib.request
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+
+AUTH_URL = "https://api.schwabapi.com/v1/oauth/authorize"
+TOKEN_URL = "https://api.schwabapi.com/v1/oauth/token"
+TRADER_BASE_URL = "https://api.schwabapi.com/trader/v1"
+DEFAULT_TOKEN_FILE = "~/.config/zenbu/schwab_tokens.json"
+
+
+class SchwabError(RuntimeError):
+    pass
+
+
+@dataclass(frozen=True)
+class SchwabConfig:
+    app_key: str
+    app_secret: str
+    redirect_uri: str
+    token_file: Path
+
+    @classmethod
+    def from_env(cls) -> "SchwabConfig":
+        app_key = os.environ.get("SCHWAB_APP_KEY", "").strip()
+        app_secret = os.environ.get("SCHWAB_APP_SECRET", "").strip()
+        redirect_uri = os.environ.get("SCHWAB_REDIRECT_URI", "").strip()
+        token_file = Path(os.environ.get("SCHWAB_TOKEN_FILE", DEFAULT_TOKEN_FILE)).expanduser()
+
+        missing = [
+            name
+            for name, value in (
+                ("SCHWAB_APP_KEY", app_key),
+                ("SCHWAB_APP_SECRET", app_secret),
+                ("SCHWAB_REDIRECT_URI", redirect_uri),
+            )
+            if not value
+        ]
+        if missing:
+            raise SchwabError("Missing required environment variables: " + ", ".join(missing))
+
+        return cls(
+            app_key=app_key,
+            app_secret=app_secret,
+            redirect_uri=redirect_uri,
+            token_file=token_file,
+        )
+
+
+@dataclass(frozen=True)
+class ApiResponse:
+    status: int
+    headers: dict[str, str]
+    body: Any
+    raw_body: str
+
+
+def build_authorization_url(app_key: str, redirect_uri: str, state: str | None = None) -> str:
+    params = {
+        "response_type": "code",
+        "client_id": app_key,
+        "redirect_uri": redirect_uri,
+    }
+    if state:
+        params["state"] = state
+    return AUTH_URL + "?" + urllib.parse.urlencode(params)
+
+
+def extract_authorization_code(code_or_url: str) -> str:
+    value = code_or_url.strip()
+    if not value:
+        raise SchwabError("Authorization code is empty")
+
+    parsed = urllib.parse.urlparse(value)
+    if parsed.scheme and parsed.netloc:
+        query = urllib.parse.parse_qs(parsed.query)
+        codes = query.get("code")
+        if not codes or not codes[0]:
+            raise SchwabError("No code= parameter found in callback URL")
+        return codes[0]
+
+    return urllib.parse.unquote(value)
+
+
+def exchange_code_for_tokens(config: SchwabConfig, code_or_url: str) -> dict[str, Any]:
+    code = extract_authorization_code(code_or_url)
+    return _token_request(
+        config,
+        {
+            "grant_type": "authorization_code",
+            "code": code,
+            "redirect_uri": config.redirect_uri,
+        },
+    )
+
+
+def refresh_tokens(config: SchwabConfig, refresh_token: str | None = None) -> dict[str, Any]:
+    token_value = refresh_token
+    if token_value is None:
+        existing = load_tokens(config.token_file)
+        token_value = existing.get("refresh_token")
+    if not token_value:
+        raise SchwabError("No refresh token available")
+
+    return _token_request(
+        config,
+        {
+            "grant_type": "refresh_token",
+            "refresh_token": token_value,
+        },
+    )
+
+
+def save_tokens(path: Path, tokens: dict[str, Any]) -> None:
+    payload = dict(tokens)
+    payload["saved_at"] = int(time.time())
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
+    path.chmod(stat.S_IRUSR | stat.S_IWUSR)
+
+
+def load_tokens(path: Path) -> dict[str, Any]:
+    if not path.exists():
+        raise SchwabError(f"Token file does not exist: {path}")
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def get_account_numbers(access_token: str) -> ApiResponse:
+    return _api_request("GET", "/accounts/accountNumbers", access_token)
+
+
+def get_accounts(access_token: str, fields: str | None = None) -> ApiResponse:
+    query = ""
+    if fields:
+        query = "?" + urllib.parse.urlencode({"fields": fields})
+    return _api_request("GET", "/accounts" + query, access_token)
+
+
+def get_account(access_token: str, account_hash: str, fields: str | None = None) -> ApiResponse:
+    query = ""
+    if fields:
+        query = "?" + urllib.parse.urlencode({"fields": fields})
+    return _api_request("GET", f"/accounts/{urllib.parse.quote(account_hash)}{query}", access_token)
+
+
+def build_equity_order(
+    action: str,
+    symbol: str,
+    quantity: float,
+    order_type: str = "MARKET",
+    price: float | None = None,
+    duration: str = "DAY",
+    session: str = "NORMAL",
+) -> dict[str, Any]:
+    normalized_action = action.upper()
+    normalized_symbol = symbol.upper()
+    normalized_order_type = order_type.upper()
+    normalized_duration = duration.upper()
+    normalized_session = session.upper()
+
+    if normalized_action not in {"BUY", "SELL"}:
+        raise SchwabError("action must be BUY or SELL")
+    if not normalized_symbol:
+        raise SchwabError("symbol is required")
+    if quantity <= 0:
+        raise SchwabError("quantity must be greater than zero")
+    if normalized_order_type not in {"MARKET", "LIMIT"}:
+        raise SchwabError("order_type must be MARKET or LIMIT")
+    if normalized_order_type == "LIMIT" and price is None:
+        raise SchwabError("LIMIT orders require --price")
+    if normalized_order_type == "MARKET" and price is not None:
+        raise SchwabError("MARKET orders cannot include --price")
+
+    order: dict[str, Any] = {
+        "orderType": normalized_order_type,
+        "session": normalized_session,
+        "duration": normalized_duration,
+        "orderStrategyType": "SINGLE",
+        "orderLegCollection": [
+            {
+                "instruction": normalized_action,
+                "quantity": quantity,
+                "instrument": {
+                    "symbol": normalized_symbol,
+                    "assetType": "EQUITY",
+                },
+            }
+        ],
+    }
+    if price is not None:
+        order["price"] = f"{price:.2f}"
+    return order
+
+
+def place_order(access_token: str, account_hash: str, order: dict[str, Any]) -> ApiResponse:
+    path = f"/accounts/{urllib.parse.quote(account_hash)}/orders"
+    return _api_request("POST", path, access_token, order)
+
+
+def _token_request(config: SchwabConfig, form: dict[str, str]) -> dict[str, Any]:
+    credentials = f"{config.app_key}:{config.app_secret}".encode("utf-8")
+    headers = {
+        "Authorization": "Basic " + base64.b64encode(credentials).decode("ascii"),
+        "Content-Type": "application/x-www-form-urlencoded",
+        "Accept": "application/json",
+    }
+    data = urllib.parse.urlencode(form).encode("utf-8")
+    request = urllib.request.Request(TOKEN_URL, data=data, headers=headers, method="POST")
+    response = _open_request(request)
+    if not isinstance(response.body, dict):
+        raise SchwabError("Token endpoint did not return a JSON object")
+    return response.body
+
+
+def _api_request(
+    method: str,
+    path: str,
+    access_token: str,
+    body: dict[str, Any] | None = None,
+) -> ApiResponse:
+    headers = {
+        "Authorization": f"Bearer {access_token}",
+        "Accept": "application/json",
+    }
+    data = None
+    if body is not None:
+        data = json.dumps(body).encode("utf-8")
+        headers["Content-Type"] = "application/json"
+
+    request = urllib.request.Request(
+        TRADER_BASE_URL + path,
+        data=data,
+        headers=headers,
+        method=method,
+    )
+    return _open_request(request)
+
+
+def _open_request(request: urllib.request.Request) -> ApiResponse:
+    try:
+        with urllib.request.urlopen(request, timeout=30) as response:
+            raw = response.read().decode("utf-8")
+            return ApiResponse(
+                status=response.status,
+                headers=dict(response.headers.items()),
+                body=_parse_json_or_text(raw),
+                raw_body=raw,
+            )
+    except urllib.error.HTTPError as error:
+        raw = error.read().decode("utf-8", errors="replace")
+        raise SchwabError(
+            f"Schwab API request failed with HTTP {error.code}: {raw or error.reason}"
+        ) from error
+    except urllib.error.URLError as error:
+        raise SchwabError(f"Schwab API request failed: {error.reason}") from error
+
+
+def _parse_json_or_text(raw: str) -> Any:
+    if not raw:
+        return None
+    try:
+        return json.loads(raw)
+    except json.JSONDecodeError:
+        return raw
+