Stonks/overview/brokerages/rest.py

198 lines
7.9 KiB
Python

import json
from datetime import datetime, timezone
from urllib.parse import urljoin
from urllib.request import Request, urlopen
from overview.brokerage_types import (
BrokerageAccountPayload,
BrokerageAssetPayload,
BrokerageHoldingPayload,
BrokerageSnapshot,
BrokerageTransactionPayload,
)
from .base import BaseBrokerageAdapter, BrokerageAdapterError
class GenericRestBrokerageAdapter(BaseBrokerageAdapter):
provider = "generic-rest"
def fetch_snapshot(self) -> BrokerageSnapshot:
accounts = self._fetch_collection(
self.connection.get_setting("account_endpoint", "/accounts")
)
assets = self._fetch_collection(
self.connection.get_setting("asset_endpoint", "/assets")
)
holdings = self._fetch_collection(
self.connection.get_setting("holding_endpoint", "/holdings")
)
transactions = self._fetch_collection(
self.connection.get_setting("transaction_endpoint", "/transactions")
)
return BrokerageSnapshot(
accounts=[self._build_account(item) for item in accounts],
assets=[self._build_asset(item) for item in assets],
holdings=[self._build_holding(item) for item in holdings],
transactions=[self._build_transaction(item) for item in transactions],
)
def _fetch_collection(self, endpoint: str) -> list[dict]:
url = urljoin(self.connection.base_url, endpoint)
request = Request(url, headers=self._headers())
try:
with urlopen(request, timeout=30) as response:
payload = json.load(response)
except Exception as exc:
raise BrokerageAdapterError(
f"Unable to fetch brokerage data from {url}: {exc}"
) from exc
if isinstance(payload, dict):
for key in ("results", "data", "items"):
value = payload.get(key)
if isinstance(value, list):
return value
if not isinstance(payload, list):
raise BrokerageAdapterError(
f"Expected a list payload from {url}, received {type(payload).__name__}"
)
return payload
def _headers(self) -> dict[str, str]:
headers = {"Accept": "application/json"}
token = self.connection.get_credential("access_token")
api_key = self.connection.get_credential("api_key")
api_key_header = self.connection.get_setting("api_key_header", "X-API-Key")
if token:
headers["Authorization"] = f"Bearer {token}"
if api_key:
headers[api_key_header] = api_key
extra_headers = self.connection.get_setting("headers", {})
if isinstance(extra_headers, dict):
headers.update({str(key): str(value) for key, value in extra_headers.items()})
return headers
def _build_account(self, item: dict) -> BrokerageAccountPayload:
account_id_field = self.connection.get_setting("account_id_field", "id")
account_name_field = self.connection.get_setting("account_name_field", "name")
account_number_field = self.connection.get_setting("account_number_field", "number")
account_balance_field = self.connection.get_setting("account_balance_field", "balance")
return BrokerageAccountPayload(
external_id=str(item[account_id_field]),
name=item[account_name_field],
number=item.get(account_number_field, ""),
balance=int(item.get(account_balance_field, 0)),
)
def _build_asset(self, item: dict) -> BrokerageAssetPayload:
asset_id_field = self.connection.get_setting("asset_id_field", "id")
asset_ticker_field = self.connection.get_setting("asset_ticker_field", "ticker")
asset_price_field = self.connection.get_setting("asset_price_field", "price")
return BrokerageAssetPayload(
external_id=str(item[asset_id_field]),
ticker=item[asset_ticker_field],
price=float(item.get(asset_price_field, 0)),
)
def _build_holding(self, item: dict) -> BrokerageHoldingPayload:
holding_id_field = self.connection.get_setting("holding_id_field", "id")
holding_account_field = self.connection.get_setting("holding_account_field", "account_id")
holding_asset_field = self.connection.get_setting("holding_asset_field", "asset_id")
holding_ticker_field = self.connection.get_setting("holding_ticker_field", "ticker")
holding_amount_field = self.connection.get_setting("holding_amount_field", "amount")
holding_cost_basis_field = self.connection.get_setting(
"holding_cost_basis_field", "cost_basis"
)
holding_purchase_date_field = self.connection.get_setting(
"holding_purchase_date_field", "purchase_date"
)
purchase_date = item.get(holding_purchase_date_field)
if purchase_date:
purchase_date = datetime.fromisoformat(
str(purchase_date).replace("Z", "+00:00")
)
else:
purchase_date = datetime.now(timezone.utc)
return BrokerageHoldingPayload(
external_id=str(item[holding_id_field]),
account_external_id=str(item[holding_account_field]),
asset_external_id=str(item[holding_asset_field]),
ticker=item[holding_ticker_field],
amount=float(item.get(holding_amount_field, 0)),
cost_basis=float(item.get(holding_cost_basis_field, 0)),
purchase_date=purchase_date,
)
def _build_transaction(self, item: dict) -> BrokerageTransactionPayload:
transaction_id_field = self.connection.get_setting("transaction_id_field", "id")
transaction_account_field = self.connection.get_setting(
"transaction_account_field",
"account_id",
)
transaction_asset_field = self.connection.get_setting(
"transaction_asset_field",
"asset_id",
)
transaction_type_field = self.connection.get_setting(
"transaction_type_field",
"transaction_type",
)
transaction_date_field = self.connection.get_setting(
"transaction_date_field",
"date",
)
transaction_value_field = self.connection.get_setting(
"transaction_value_field",
"value",
)
transaction_name_field = self.connection.get_setting(
"transaction_name_field",
"name",
)
transaction_quantity_field = self.connection.get_setting(
"transaction_quantity_field",
"quantity",
)
transaction_realized_gain_field = self.connection.get_setting(
"transaction_realized_gain_field",
"realized_gain",
)
transaction_date = item.get(transaction_date_field)
if transaction_date:
transaction_date = datetime.fromisoformat(
str(transaction_date).replace("Z", "+00:00")
)
else:
transaction_date = datetime.now(timezone.utc)
asset_external_id = item.get(transaction_asset_field)
if asset_external_id is not None:
asset_external_id = str(asset_external_id)
quantity = item.get(transaction_quantity_field)
realized_gain = item.get(transaction_realized_gain_field)
return BrokerageTransactionPayload(
external_id=str(item[transaction_id_field]),
account_external_id=str(item[transaction_account_field]),
asset_external_id=asset_external_id,
transaction_type=str(item[transaction_type_field]),
date=transaction_date,
value=float(item.get(transaction_value_field, 0)),
name=str(item.get(transaction_name_field, "")),
quantity=float(quantity) if quantity is not None else None,
realized_gain=float(realized_gain) if realized_gain is not None else None,
)