265 lines
9.2 KiB
Python
265 lines
9.2 KiB
Python
import httpx
|
||
from httpx import BasicAuth, Timeout
|
||
import pandas as pd
|
||
from tqdm import tqdm
|
||
from pathlib import Path
|
||
import logging
|
||
import time
|
||
import random
|
||
|
||
|
||
class SimpleAlphaFetcher:
|
||
def __init__(self, base_path=None):
|
||
"""
|
||
初始化 Alpha 获取器
|
||
|
||
Args:
|
||
base_path: 输出文件保存路径
|
||
"""
|
||
self.client = None
|
||
self.base_path = Path(base_path) if base_path else Path.cwd()
|
||
self.logger = self._setup_logger()
|
||
|
||
# 登录
|
||
self.login()
|
||
|
||
def _setup_logger(self):
|
||
"""设置日志记录器"""
|
||
logger = logging.getLogger(__name__)
|
||
logger.setLevel(logging.INFO)
|
||
|
||
if not logger.handlers:
|
||
handler = logging.StreamHandler()
|
||
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
|
||
handler.setFormatter(formatter)
|
||
logger.addHandler(handler)
|
||
|
||
return logger
|
||
|
||
def login(self):
|
||
"""登录 WorldQuant Brain API"""
|
||
try:
|
||
# 从 nacos 获取账号密码
|
||
with httpx.Client(timeout=10.0) as temp_client:
|
||
nacos_resp = temp_client.get(
|
||
'http://192.168.31.41:30848/nacos/v1/cs/configs?dataId=wq_account&group=quantify'
|
||
)
|
||
|
||
if nacos_resp.status_code != 200:
|
||
self.logger.error('获取账号密码失败')
|
||
return False
|
||
|
||
config = nacos_resp.json()
|
||
username = config.get('user_name')
|
||
password = config.get('password')
|
||
|
||
if not username or not password:
|
||
self.logger.error('账号密码不完整')
|
||
return False
|
||
|
||
self.logger.info(f"正在登录账户: {username}")
|
||
|
||
# 创建客户端并设置超时
|
||
timeout = Timeout(connect=30.0, read=60.0, write=30.0, pool=30.0)
|
||
self.client = httpx.Client(
|
||
auth=BasicAuth(username, password),
|
||
timeout=timeout
|
||
)
|
||
|
||
# 发送登录请求
|
||
response = self.client.post('https://api.worldquantbrain.com/authentication')
|
||
|
||
if response.status_code == 201:
|
||
self.logger.info("登录成功!")
|
||
return True
|
||
else:
|
||
self.logger.error(f"登录失败: {response.status_code} - {response.text}")
|
||
self.client.close()
|
||
self.client = None
|
||
return False
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"登录异常: {e}")
|
||
return False
|
||
|
||
def fetch_alphas(
|
||
self,
|
||
max_pages: int = 100,
|
||
limit: int = 100,
|
||
delay: int = 1,
|
||
region: str = "USA",
|
||
universe: str = "TOP3000",
|
||
hidden: str = "false",
|
||
output_file_name: str = "alpha_list.csv",
|
||
mode: str = "w",
|
||
max_retries: int = 3,
|
||
):
|
||
"""
|
||
获取 Alpha 列表
|
||
|
||
Args:
|
||
max_pages: 最大爬取页数
|
||
limit: 每页数量
|
||
delay: API 请求的延迟设置
|
||
region: 市场区域
|
||
universe: 股票池
|
||
hidden: 是否搜索已隐藏的 Alpha
|
||
output_file_name: 输出的 CSV 文件名
|
||
mode: 写入模式："w"(覆盖写入) 或 "a"(追加写入)
|
||
max_retries: 请求失败时的最大重试次数
|
||
|
||
Returns:
|
||
list: Alpha 列表
|
||
"""
|
||
if not self.client:
|
||
self.logger.error("客户端未登录，无法执行获取")
|
||
return []
|
||
|
||
fetched_alphas = []
|
||
offset = 0
|
||
total_accessed = 0
|
||
pages_fetched = 0
|
||
|
||
# 先获取总数
|
||
count_url = (
|
||
f"https://api.worldquantbrain.com/users/self/alphas?stage=IS&hidden={hidden}"
|
||
f"&limit=1&settings.delay={delay}&settings.region={region}&status=UNSUBMITTED%1FIS_FAIL&settings.universe={universe}"
|
||
)
|
||
|
||
total_available = 0
|
||
for attempt in range(max_retries):
|
||
try:
|
||
count_response = self.client.get(count_url)
|
||
total_available = count_response.json()["count"]
|
||
break
|
||
except Exception as e:
|
||
self.logger.warning(f"获取 Alpha 总数失败 (尝试 {attempt+1}/{max_retries}): {e}")
|
||
if attempt < max_retries - 1:
|
||
time.sleep(random.uniform(3, 5))
|
||
else:
|
||
self.logger.error("获取 Alpha 总数最终失败")
|
||
return []
|
||
|
||
if total_available == 0:
|
||
self.logger.warning("未找到任何 Alpha")
|
||
return []
|
||
|
||
# 计算实际最大页数
|
||
actual_max_pages = min(max_pages, (total_available + limit - 1) // limit)
|
||
self.logger.info(f"共找到 {total_available} 个 Alpha，计划爬取 {actual_max_pages} 页...")
|
||
|
||
pbar = tqdm(total=actual_max_pages, desc="爬取 Alpha 页面", unit="页")
|
||
|
||
while pages_fetched < actual_max_pages:
|
||
# 构建请求 URL
|
||
url = (
|
||
f"https://api.worldquantbrain.com/users/self/alphas?stage=IS&limit={limit}"
|
||
f"&offset={offset}&settings.delay={delay}&settings.region={region}&hidden={hidden}&status=UNSUBMITTED%1FIS_FAIL&settings.universe={universe}"
|
||
)
|
||
|
||
try:
|
||
# 使用重试机制
|
||
response = None
|
||
for attempt in range(max_retries):
|
||
try:
|
||
response = self.client.get(url)
|
||
break
|
||
except Exception as e:
|
||
self.logger.warning(f"请求失败 (尝试 {attempt+1}/{max_retries}): {e}")
|
||
if attempt < max_retries - 1:
|
||
time.sleep(random.uniform(3, 5))
|
||
else:
|
||
raise
|
||
|
||
if response.status_code == 400:
|
||
self.logger.warning(f"遇到 API 限制 (offset={offset})，停止获取")
|
||
break
|
||
|
||
response_data = response.json()
|
||
|
||
if not isinstance(response_data, dict) or "results" not in response_data:
|
||
self.logger.error(f"API 返回了意外的数据: {response_data}")
|
||
break
|
||
|
||
alphas = response_data["results"]
|
||
|
||
if not alphas:
|
||
break
|
||
|
||
fetched_alphas.extend(alphas)
|
||
total_accessed += len(alphas)
|
||
pages_fetched += 1
|
||
|
||
# 更新进度条
|
||
pbar.update(1)
|
||
pbar.set_postfix({
|
||
"Region": region,
|
||
"已获取": total_accessed,
|
||
"本页": len(alphas),
|
||
})
|
||
|
||
if len(alphas) < limit:
|
||
self.logger.info(f"本页数据不足 {limit} 条，已到达末尾")
|
||
break
|
||
|
||
offset += limit
|
||
|
||
# 添加延迟避免请求过快
|
||
time.sleep(0.5)
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"请求失败: {e}")
|
||
break
|
||
|
||
pbar.close()
|
||
|
||
if not fetched_alphas:
|
||
self.logger.warning("未获取到任何 Alpha！")
|
||
return []
|
||
|
||
# 按 fitness 排序
|
||
df = pd.DataFrame(fetched_alphas)
|
||
df["temp_fitness"] = df.apply(
|
||
lambda row: row["is"].get("fitness", 0) if isinstance(row.get("is"), dict) else 0,
|
||
axis=1
|
||
)
|
||
df_sorted = df.sort_values(by="temp_fitness", ascending=False)
|
||
df_sorted = df_sorted.drop("temp_fitness", axis=1)
|
||
|
||
output_path = self.base_path / output_file_name
|
||
|
||
if mode == "w":
|
||
df_sorted.to_csv(output_path, index=False)
|
||
elif mode == "a":
|
||
df_sorted.to_csv(output_path, mode="a", index=False, header=False)
|
||
|
||
self.logger.info(f"Alpha 列表已保存！共 {len(fetched_alphas)} 条记录！\n文件路径: {output_path}")
|
||
|
||
return fetched_alphas
|
||
|
||
def close(self):
|
||
"""关闭客户端"""
|
||
if self.client:
|
||
self.client.close()
|
||
|
||
|
||
# 使用示例
|
||
if __name__ == "__main__":
|
||
# 创建获取器实例（会自动登录）
|
||
fetcher = SimpleAlphaFetcher()
|
||
|
||
if fetcher.client:
|
||
# 获取 Alpha 列表
|
||
results = fetcher.fetch_alphas(
|
||
max_pages=100, # 最多爬取 100 页
|
||
limit=100, # 每页 100 条
|
||
region="USA", # 美股市场
|
||
universe="TOP3000", # TOP3000 股票池
|
||
output_file_name="alpha_list.csv" # 输出文件名
|
||
)
|
||
|
||
print(f"成功获取 {len(results) if results else 0} 个 Alpha")
|
||
|
||
# 关闭客户端
|
||
fetcher.close()
|