++

wqb-get-alphas/simple_alpha_fetcher.py

@@ -0,0 +1,264 @@
+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()