alpha_tools/wqb-get-alphas/simple_alpha_fetcher.py

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()