yuxuanhui
526 lines
24 KiB
Python
526 lines
24 KiB
Python
"""WorldQuant adapter. Explicit descriptions allow PATCH; Alpha submission is not exposed.
|
||
|
||
No upstream response body or request headers are included in exceptions: they may
|
||
contain credentials, cookies, or temporary authentication links.
|
||
"""
|
||
|
||
import asyncio
|
||
import math
|
||
import random
|
||
import re
|
||
from contextvars import ContextVar
|
||
from datetime import datetime, timedelta, timezone
|
||
from email.utils import parsedate_to_datetime
|
||
from typing import Awaitable, Callable
|
||
from urllib.parse import urljoin, urlparse
|
||
|
||
import httpx
|
||
|
||
|
||
class WqError(Exception):
|
||
def __init__(self, message: str, code: str = "upstream_error"):
|
||
super().__init__(message)
|
||
self.code = code
|
||
|
||
|
||
class VerificationRequired(WqError):
|
||
def __init__(self, url: str):
|
||
super().__init__("请在 WorldQuant 完成人工验证，再点击继续验证", "verification_required")
|
||
self.url = url
|
||
|
||
|
||
class SimulationDeferred(WqError):
|
||
def __init__(self, message, delay=5, code="rate_limited"):
|
||
super().__init__(message, code)
|
||
self.delay = delay
|
||
|
||
|
||
class WqClient:
|
||
def __init__(self, settings, transport=None, sleep=asyncio.sleep):
|
||
self.settings = settings
|
||
self.client = httpx.AsyncClient(
|
||
base_url=settings.wq_base_url,
|
||
timeout=settings.request_timeout,
|
||
transport=transport,
|
||
follow_redirects=False,
|
||
)
|
||
self.lock = asyncio.Lock()
|
||
self.authenticated = False
|
||
self.auth_generation = 0
|
||
self.credentials: tuple[str, str] | None = None
|
||
self.verification_url: str | None = None
|
||
self.permissions: list[str] | None = None
|
||
self.session_expires_at: datetime | None = None
|
||
self.session_duration: float | None = None
|
||
self.sleep = sleep
|
||
self._retry_hook = ContextVar("wq_retry_hook", default=None)
|
||
|
||
@property
|
||
def on_retry(self) -> Callable[[float], Awaitable[None]] | None:
|
||
return self._retry_hook.get()
|
||
|
||
@on_retry.setter
|
||
def on_retry(self, value):
|
||
# Sync and simulation tasks share a session, never each other's retry callback.
|
||
self._retry_hook.set(value)
|
||
|
||
def simulation_url(self, value):
|
||
"""Accept only same-origin simulation resources; never forward cookies elsewhere."""
|
||
base = urlparse(self.settings.wq_base_url)
|
||
url = urlparse(urljoin(self.settings.wq_base_url, value))
|
||
if (
|
||
url.scheme != base.scheme
|
||
or url.netloc != base.netloc
|
||
or url.query
|
||
or url.fragment
|
||
or not re.fullmatch(r"/simulations/[A-Za-z0-9_-]+", url.path)
|
||
):
|
||
raise WqError("模拟引用地址无法确认", "invalid_simulation_url")
|
||
return url.geturl()
|
||
|
||
async def submit_simulations(self, payload):
|
||
"""One POST only. Transport/5xx/invalid acknowledgement may already be accepted."""
|
||
try:
|
||
response = await self.client.post(
|
||
"/simulations", json=payload[0] if len(payload) == 1 else payload
|
||
)
|
||
except httpx.TransportError:
|
||
raise WqError("提交结果未知，禁止自动重提，请核对平台任务", "submission_unknown") from None
|
||
if response.status_code == 429:
|
||
raise SimulationDeferred(
|
||
"平台限流，暂停后续提交", self.retry_delay(response.headers.get("Retry-After"), 0)
|
||
)
|
||
if response.status_code == 401:
|
||
self.authenticated = False
|
||
raise SimulationDeferred("平台会话过期，重新认证后继续", 2, "session_expired")
|
||
if response.status_code in (400, 403, 404, 422):
|
||
raise WqError(f"平台拒绝回测提交（HTTP {response.status_code}）", "submission_rejected")
|
||
if response.status_code != 201 or not response.headers.get("Location"):
|
||
raise WqError("平台未返回可靠提交凭证，请核对后再处理", "submission_unknown")
|
||
try:
|
||
return self.simulation_url(response.headers["Location"])
|
||
except WqError:
|
||
raise WqError("平台已响应但模拟引用无法确认，禁止自动重提", "submission_unknown") from None
|
||
|
||
async def poll_simulation(self, url):
|
||
response = await self._request("GET", self.simulation_url(url))
|
||
if response.status_code == 401:
|
||
self.authenticated = False
|
||
raise SimulationDeferred("平台会话过期，重新认证后继续", 2, "session_expired")
|
||
if response.status_code not in (200, 202):
|
||
raise WqError(
|
||
f"模拟查询失败（HTTP {response.status_code}），保留原任务", "simulation_unavailable"
|
||
)
|
||
try:
|
||
data = response.json()
|
||
if not isinstance(data, dict):
|
||
raise ValueError()
|
||
except ValueError:
|
||
raise WqError("模拟响应格式无法识别，保留原任务", "invalid_response") from None
|
||
return data, self.retry_delay(response.headers["Retry-After"], 0) if response.headers.get(
|
||
"Retry-After"
|
||
) else 0
|
||
|
||
async def close(self):
|
||
await self.client.aclose()
|
||
|
||
def disconnect(self):
|
||
self.authenticated = False
|
||
self.credentials = None
|
||
self.verification_url = None
|
||
self.permissions = None
|
||
self.session_expires_at = None
|
||
self.session_duration = None
|
||
self.client.cookies.clear()
|
||
|
||
def capture_session(self, response: httpx.Response):
|
||
"""Keep only permissions and expiry metadata; never retain the authentication body."""
|
||
self.permissions, self.session_expires_at, self.session_duration = None, None, None
|
||
try:
|
||
data = response.json()
|
||
except ValueError:
|
||
return
|
||
if not isinstance(data, dict):
|
||
return
|
||
permissions = data.get("permissions")
|
||
if isinstance(permissions, list) and all(isinstance(p, str) for p in permissions):
|
||
self.permissions = list(dict.fromkeys(permissions))
|
||
token = data.get("token")
|
||
expiry = token.get("expiry") if isinstance(token, dict) else None
|
||
if (
|
||
isinstance(expiry, (int, float))
|
||
and not isinstance(expiry, bool)
|
||
and math.isfinite(expiry)
|
||
and 0 <= expiry <= 31536000
|
||
):
|
||
self.session_duration = expiry
|
||
self.session_expires_at = datetime.now(timezone.utc) + timedelta(seconds=expiry)
|
||
|
||
def session_info(self):
|
||
remaining = (
|
||
max(0, int((self.session_expires_at - datetime.now(timezone.utc)).total_seconds()))
|
||
if self.session_expires_at
|
||
else None
|
||
)
|
||
return {
|
||
"authenticated": self.authenticated and remaining != 0,
|
||
"expires_at": self.session_expires_at,
|
||
"remaining_seconds": remaining,
|
||
"total_seconds": self.session_duration,
|
||
}
|
||
|
||
def safe_verification_url(self, response: httpx.Response) -> str:
|
||
url = urljoin(str(response.url), response.headers.get("Location", ""))
|
||
base = urlparse(self.settings.wq_base_url)
|
||
parsed = urlparse(url)
|
||
if (
|
||
not response.headers.get("Location")
|
||
or parsed.scheme != base.scheme
|
||
or parsed.netloc != base.netloc
|
||
):
|
||
raise WqError("平台返回了无法确认的验证地址，请到 BRAIN 官网检查账户", "invalid_verification")
|
||
return url
|
||
|
||
@staticmethod
|
||
def retry_delay(value: str | None, attempt: int) -> float:
|
||
if value:
|
||
try:
|
||
seconds = float(value)
|
||
except ValueError:
|
||
try:
|
||
seconds = (parsedate_to_datetime(value) - datetime.now(timezone.utc)).total_seconds()
|
||
except (TypeError, ValueError):
|
||
seconds = -1
|
||
if math.isfinite(seconds) and seconds >= 0:
|
||
return seconds
|
||
return min(2**attempt + random.uniform(0, 0.5), 30)
|
||
|
||
async def _request(self, method: str, path: str, **kwargs) -> httpx.Response:
|
||
for attempt in range(self.settings.retry_attempts):
|
||
try:
|
||
response = await self.client.request(method, path, **kwargs)
|
||
except (httpx.TimeoutException, httpx.NetworkError):
|
||
if attempt + 1 == self.settings.retry_attempts:
|
||
raise WqError("WorldQuant 网络请求失败，请稍后重试", "network_error") from None
|
||
response = None
|
||
if response is not None and response.status_code not in (429, 500, 502, 503, 504):
|
||
return response
|
||
if attempt + 1 == self.settings.retry_attempts:
|
||
raise WqError("WorldQuant 暂时限流或不可用，重试预算已用完", "retry_exhausted")
|
||
delay = self.retry_delay(
|
||
response.headers.get("Retry-After") if response is not None else None, attempt
|
||
)
|
||
if self.on_retry:
|
||
await self.on_retry(delay)
|
||
await self.sleep(delay)
|
||
raise WqError("重试失败")
|
||
|
||
async def authenticate(self, email: str, password: str, force=False, stale_generation=None):
|
||
async with self.lock:
|
||
if (
|
||
self.authenticated
|
||
and not force
|
||
and (stale_generation is None or stale_generation != self.auth_generation)
|
||
):
|
||
return
|
||
if self.verification_url and not force:
|
||
raise VerificationRequired(self.verification_url)
|
||
self.credentials = (email, password)
|
||
self.authenticated = False
|
||
self.permissions, self.session_expires_at, self.session_duration = None, None, None
|
||
if force:
|
||
self.client.cookies.clear()
|
||
self.verification_url = None
|
||
response = await self._request("POST", "/authentication", auth=httpx.BasicAuth(email, password))
|
||
if (
|
||
response.status_code == 401
|
||
and response.headers.get("WWW-Authenticate", "").lower() == "persona"
|
||
):
|
||
self.verification_url = self.safe_verification_url(response)
|
||
raise VerificationRequired(self.verification_url)
|
||
if response.status_code in (401, 403):
|
||
raise WqError("平台认证未通过，请检查凭据或账户访问权限", "authentication_failed")
|
||
if response.status_code not in (200, 201):
|
||
raise WqError(f"平台认证返回异常状态 {response.status_code}", "authentication_failed")
|
||
self.authenticated = True
|
||
self.capture_session(response)
|
||
self.auth_generation += 1
|
||
self.verification_url = None
|
||
|
||
async def verify(self):
|
||
"""Continue the same server-side challenge session after the human visits the URL."""
|
||
if not self.verification_url:
|
||
if not self.credentials:
|
||
raise WqError("验证会话已结束，请重新连接", "authentication_failed")
|
||
await self.authenticate(*self.credentials, force=True)
|
||
return
|
||
response = await self._request("POST", self.verification_url)
|
||
if response.status_code in (200, 201):
|
||
self.authenticated = True
|
||
self.capture_session(response)
|
||
self.auth_generation += 1
|
||
self.verification_url = None
|
||
elif response.status_code in (401, 403, 202):
|
||
raise VerificationRequired(self.verification_url)
|
||
else:
|
||
raise WqError("验证会话已失效，请重新连接", "authentication_failed")
|
||
|
||
async def get(self, path: str, params=None, headers=None):
|
||
return await self._read_json("GET", path, params=params, headers=headers)
|
||
|
||
async def _read_json(self, method: str, path: str, *, allow_list=False, poll_attempts=None, wait_for_retry_header=False, **kwargs):
|
||
"""Authenticated read with shared refresh/retry handling; callers use GET or OPTIONS."""
|
||
if not self.credentials:
|
||
raise WqError("请先连接 WorldQuant", "disconnected")
|
||
if not self.authenticated:
|
||
await self.authenticate(*self.credentials)
|
||
refreshed = False
|
||
attempts = self.settings.retry_attempts if poll_attempts is None else poll_attempts
|
||
for attempt in range(attempts):
|
||
generation = self.auth_generation
|
||
response = await self._request(method, path, **kwargs)
|
||
if response.status_code == 401 and not refreshed:
|
||
await self.authenticate(*self.credentials, stale_generation=generation)
|
||
refreshed = True
|
||
continue
|
||
if response.status_code in (401, 403):
|
||
raise WqError("无权访问该平台资源", "access_denied")
|
||
if response.status_code == 404:
|
||
raise WqError("平台资源不存在或不可访问", "not_found")
|
||
if response.status_code not in (200, 201, 202):
|
||
raise WqError(f"WorldQuant 请求返回状态 {response.status_code}")
|
||
# Recordsets may return 200/202 with Retry-After before results exist.
|
||
if (
|
||
response.headers.get("Retry-After")
|
||
and (wait_for_retry_header or self.retry_delay(response.headers["Retry-After"], 0) > 0)
|
||
):
|
||
if attempt + 1 == attempts:
|
||
raise WqError("平台数据仍在准备，请稍后重试", "pending")
|
||
delay = self.retry_delay(response.headers["Retry-After"], attempt)
|
||
if self.on_retry:
|
||
await self.on_retry(delay)
|
||
await self.sleep(delay)
|
||
continue
|
||
try:
|
||
result = response.json()
|
||
if not isinstance(result, dict) and not (allow_list and isinstance(result, list)):
|
||
raise ValueError()
|
||
return result
|
||
except ValueError:
|
||
raise WqError("平台返回的数据格式无法识别", "invalid_response") from None
|
||
raise WqError("平台数据尚未就绪", "pending")
|
||
|
||
async def profile(self):
|
||
return await self.get("/users/self")
|
||
|
||
async def get_pyramid_multipliers(self):
|
||
"""Return current Pyramid multipliers using the shared authenticated session."""
|
||
return await self.get("/users/self/activities/pyramid-multipliers")
|
||
|
||
async def get_pyramid_alphas(self, start_date=None, end_date=None):
|
||
"""Return distribution for optional ISO dates; only missing routes trigger fallback.
|
||
|
||
Paths mirror the cnhk adapter. Access denial, authentication and network
|
||
failures must retain their meaning instead of trying unrelated routes.
|
||
"""
|
||
params = {}
|
||
if start_date is not None:
|
||
params["startDate"] = start_date
|
||
if end_date is not None:
|
||
params["endDate"] = end_date
|
||
for path in (
|
||
"/users/self/activities/pyramid-alphas",
|
||
"/users/self/pyramid/alphas",
|
||
"/activities/pyramid-alphas",
|
||
):
|
||
try:
|
||
return await self.get(path, params=params)
|
||
except WqError as exc:
|
||
if exc.code != "not_found":
|
||
raise
|
||
raise WqError("当前账户的 Pyramid 分布接口不可用", "not_found")
|
||
|
||
async def account_usage(self):
|
||
"""Read independent account resources; unavailable sections do not hide the profile."""
|
||
from .account_data import usage_snapshot
|
||
|
||
resources = {
|
||
"simulations": "/users/self/activities/simulations",
|
||
"submissions": "/users/self/activities/submissions",
|
||
"alphas": "/users/self/alphas/summary",
|
||
}
|
||
data, errors = {}, {}
|
||
for key, path in resources.items():
|
||
try:
|
||
data[key] = await self.get(
|
||
path, headers={"Accept": "application/json;version=4.0"} if key == "alphas" else None
|
||
)
|
||
except VerificationRequired:
|
||
raise
|
||
except WqError as exc:
|
||
if exc.code in ("authentication_failed", "disconnected"):
|
||
raise
|
||
errors[key] = str(exc)
|
||
return usage_snapshot(data, errors)
|
||
|
||
async def alphas(self, submission, hidden, offset, before, *, date_from=None, date_to=None):
|
||
# Cover every platform stage; submitted records are not assumed to be OS only.
|
||
status_key = "status" if submission == "UNSUBMITTED" else "status!"
|
||
params = {
|
||
status_key: "UNSUBMITTED",
|
||
"hidden": str(hidden).lower(),
|
||
"limit": 100,
|
||
"offset": offset,
|
||
"order": "dateCreated",
|
||
"dateCreated<": before,
|
||
}
|
||
if date_from is not None and date_to is not None:
|
||
# Daily intervals are [UTC midnight, next midnight). Submitted dates
|
||
# use the actual submission time, independently of creation or stage.
|
||
field = "dateCreated" if submission == "UNSUBMITTED" else "dateSubmitted"
|
||
# BRAIN treats the query separator '=' as part of its comparison
|
||
# syntax: a '>=' key adds an extra '=' and makes the date invalid.
|
||
# Both bounds are inclusive upstream; exclude the next UTC midnight.
|
||
upper = (datetime.fromisoformat(date_to) - timedelta(microseconds=1)).isoformat()
|
||
params.update({f"{field}>": date_from, f"{field}<": upper, "order": field})
|
||
if field == "dateCreated":
|
||
params["dateCreated<"] = min(before, upper)
|
||
return await self.get("/users/self/alphas", params)
|
||
|
||
async def alpha(self, alpha_id):
|
||
return await self.get(f"/alphas/{alpha_id}")
|
||
|
||
async def patch_descriptions(self, alpha_id, descriptions):
|
||
"""Write only reviewed descriptions; ambiguous writes require GET reconciliation.
|
||
|
||
No blind transport/5xx retry: the first PATCH may already have succeeded.
|
||
A job retry reads the current Alpha before deciding whether PATCH is needed.
|
||
"""
|
||
if not re.fullmatch(r"[A-Za-z0-9_-]{1,100}", alpha_id) or not descriptions or not set(descriptions) <= {"regular", "selection", "combo"}:
|
||
raise WqError("Description 写入参数无效", "invalid_description")
|
||
if any(not isinstance(value, dict) or set(value) != {"description"} or not isinstance(value["description"], str) for value in descriptions.values()):
|
||
raise WqError("仅允许写入 Description", "invalid_description")
|
||
if not self.credentials:
|
||
raise WqError("请先连接 WorldQuant", "disconnected")
|
||
if not self.authenticated:
|
||
await self.authenticate(*self.credentials)
|
||
for attempt in range(2):
|
||
generation = self.auth_generation
|
||
try:
|
||
response = await self.client.patch(f"/alphas/{alpha_id}", json=descriptions)
|
||
except httpx.TransportError:
|
||
raise WqError("Description 写回结果未知；重试任务将先核对平台内容", "description_unknown") from None
|
||
if response.status_code == 401 and attempt == 0:
|
||
await self.authenticate(*self.credentials, stale_generation=generation)
|
||
continue
|
||
if response.status_code in (200, 204):
|
||
return
|
||
raise WqError(f"Description 写回未确认（HTTP {response.status_code}）；可重试任务核对平台内容", "description_write_failed")
|
||
|
||
async def submission_check(self, alpha_id):
|
||
"""cnhkmcp /check contract: wait on Retry-After and return only is.checks."""
|
||
if not re.fullmatch(r"[A-Za-z0-9_-]{1,100}", alpha_id):
|
||
raise WqError("Alpha ID 无效", "invalid_alpha")
|
||
raw = await self._read_json(
|
||
"GET", f"/alphas/{alpha_id}/check", poll_attempts=self.settings.pnl_poll_attempts,
|
||
wait_for_retry_header=True,
|
||
)
|
||
metrics = raw.get("is")
|
||
checks = metrics.get("checks") if isinstance(metrics, dict) else None
|
||
if not isinstance(checks, list) or any(not isinstance(item, dict) for item in checks):
|
||
raise WqError("平台检查尚无有效 is.checks 结果，请稍后重试", "pending")
|
||
return checks
|
||
|
||
async def pnl(self, alpha_id):
|
||
"""Wait for slow PnL generation separately from transport-error retries.
|
||
|
||
Each pending response uses the platform's Retry-After delay. Exhausting
|
||
the bounded polling budget raises WqError with code 'pending'.
|
||
"""
|
||
return await self._read_json(
|
||
"GET", f"/alphas/{alpha_id}/recordsets/pnl", poll_attempts=self.settings.pnl_poll_attempts
|
||
)
|
||
|
||
async def catalog_page(self, scope, dataset_id, offset):
|
||
"""Read a single scoped page. IDs are query parameters, never upstream paths."""
|
||
params = {
|
||
"instrumentType": scope["instrument_type"],
|
||
"region": scope["region"],
|
||
"universe": scope["universe"],
|
||
"delay": scope["delay"],
|
||
"limit": 50,
|
||
"offset": offset,
|
||
}
|
||
if dataset_id is not None:
|
||
params["dataset.id"] = dataset_id
|
||
return await self.get("/data-fields" if dataset_id else "/data-sets", params)
|
||
|
||
async def operators(self, offset=0):
|
||
"""The operator endpoint has both list and paginated response forms."""
|
||
return await self._read_json("GET", "/operators", allow_list=True, params={"limit": 100, "offset": offset})
|
||
|
||
async def field_availability(self, field_id, scope):
|
||
"""Use a validated identifier, never an arbitrary upstream path or URL."""
|
||
if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", field_id):
|
||
raise WqError("字段标识格式无效", "invalid_field")
|
||
return await self.get(f"/data-fields/{field_id}", {
|
||
"instrumentType": scope.instrument_type, "region": scope.region,
|
||
"universe": scope.universe, "delay": scope.delay,
|
||
})
|
||
|
||
async def run_super_selection(self, query):
|
||
"""Read cnhk super-selection contract with bounded async retries and shared authentication."""
|
||
allowed = {"selection", "instrumentType", "region", "delay", "selectionLimit", "selectionHandling"}
|
||
if set(query) != allowed:
|
||
raise WqError("Selection 参数不完整或包含未知键", "invalid_selection")
|
||
return await self._read_json("GET", "/simulations/super-selection", params=query, allow_list=True,
|
||
wait_for_retry_header=True)
|
||
|
||
async def research_setting_options(self):
|
||
"""Snapshot full setting choices for constrained research, including neutralization."""
|
||
return await self._read_json("OPTIONS", "/simulations")
|
||
|
||
async def get_platform_setting_options(self):
|
||
"""Read platform choices for the connected account; malformed responses raise WqError."""
|
||
data = await self._read_json("OPTIONS", "/simulations")
|
||
try:
|
||
children = data["actions"]["POST"]["settings"]["children"]
|
||
|
||
def choices(key, instrument=None, region=None):
|
||
value = children[key]["choices"]
|
||
if instrument is not None:
|
||
value = value["instrumentType"][instrument]
|
||
if region is not None:
|
||
value = value["region"][region]
|
||
values = [item["value"] for item in value]
|
||
if not values:
|
||
raise ValueError()
|
||
return values
|
||
|
||
instruments = choices("instrumentType")
|
||
regions = {}
|
||
rows = []
|
||
for instrument in instruments:
|
||
regions[instrument] = choices("region", instrument)
|
||
for region in regions[instrument]:
|
||
universes = choices("universe", instrument, region)
|
||
for delay in choices("delay", instrument, region):
|
||
if type(delay) is not int or delay < 0:
|
||
raise ValueError()
|
||
if not all(
|
||
isinstance(v, str) and v and "|" not in v
|
||
for v in [instrument, region, *universes]
|
||
):
|
||
raise ValueError()
|
||
rows.append(
|
||
dict(instrument_type=instrument, region=region, delay=delay, universes=universes)
|
||
)
|
||
return dict(
|
||
instrument_options=rows,
|
||
instrument_types=instruments,
|
||
regions_by_type=regions,
|
||
total_combinations=len(rows),
|
||
)
|
||
except (KeyError, TypeError, ValueError):
|
||
raise WqError("平台配置选项格式无法识别，请稍后重试", "invalid_response") from None
|