import httpx import json from concurrent.futures import ThreadPoolExecutor, as_completed from os.path import expanduser, exists from httpx import BasicAuth from time import sleep from datetime import datetime class WorldQuantBrainAPI: def __init__(self, credentials_file='brain_credentials.txt'): """初始化API客户端""" self.credentials_file = expanduser(credentials_file) self.client = None self.brain_api_url = 'https://api.worldquantbrain.com' def login(self): """登录认证""" try: # 检查凭证文件是否存在 if not exists(self.credentials_file): print(f"❌ 凭证文件不存在: {self.credentials_file}") return False # Load credentials with open(self.credentials_file) as f: file_content = f.read().strip() with open(self.credentials_file) as f: credentials = eval(f.read()) # Extract username and password from the list if len(credentials) != 2: print(f"❌ 凭证格式错误，应该是 [username, password] 格式，但得到: {credentials}") return False username, password = credentials # Create a client with basic authentication self.client = httpx.Client(auth=BasicAuth(username, password)) # Send authentication request response = self.client.post(f'{self.brain_api_url}/authentication') print(f"认证状态: {response.status_code}") if response.status_code == 201: print("✅ 登录成功!") return True else: print(f"❌ 登录失败: {response.text}") return False except json.JSONDecodeError as e: print(f"❌ JSON解析错误: {e}") print(f"📄 文件内容: {file_content if 'file_content' in locals() else '无法读取'}") return False except Exception as e: print(f"❌ 登录过程中出错: {str(e)}") return False def batch_simulate_alphas(self, alphas_config_file='alphas.json', max_workers=3): """批量测试多个Alpha因子（使用线程池）""" # 检查配置文件是否存在 if not exists(alphas_config_file): print(f"❌ 配置文件 {alphas_config_file} 不存在") return [] # 读取Alpha配置 try: with open(alphas_config_file, 'r', encoding='utf-8') as f: file_content = f.read().strip() print(f"📄 Alpha配置文件内容预览: {file_content[:200]}...") # 只显示前200个字符 with open(alphas_config_file, 'r', encoding='utf-8') as f: alpha_configs = json.load(f) print(f"✅ 成功读取 {len(alpha_configs)} 个Alpha配置") for i, config in enumerate(alpha_configs): print(f" {i+1}. {config.get('name', '未命名')}") except json.JSONDecodeError as e: print(f"❌ 配置文件 JSON格式错误: {e}") print(f"📄 出错的文件内容: {file_content if 'file_content' in locals() else '无法读取'}") return [] except Exception as e: print(f"❌ 读取配置文件时出错: {str(e)}") return [] print(f"📁 读取到 {len(alpha_configs)} 个Alpha配置") # 登录 if not self.login(): print("❌ 登录失败，无法继续") return [] results = [] # 使用线程池并发执行 with ThreadPoolExecutor(max_workers=max_workers) as executor: # 提交所有任务 future_to_alpha = { executor.submit(self.simulate_alpha, config): config for config in alpha_configs } # 收集结果 for future in as_completed(future_to_alpha): config = future_to_alpha[future] try: result = future.result() results.append(result) except Exception as e: print(f"❌ {config.get('name', '未知Alpha')} 执行失败: {str(e)}") results.append({ 'name': config.get('name', '未知Alpha'), 'success': False, 'error': str(e) }) # 打印汇总结果 self._print_summary(results) return results def simulate_alpha(self, alpha_config): """模拟单个Alpha因子""" if self.client is None: raise Exception("请先调用 login() 方法登录") name = alpha_config.get('name', '未命名Alpha') expression = alpha_config['expression'] settings = alpha_config.get('settings', {}) print(f"\n🚀 开始模拟: {name}") print(f"📝 表达式: {expression}") simulation_data = { 'type': 'REGULAR', 'settings': settings, 'regular': expression } try: # 发送模拟数据 sim_resp = self.client.post( f'{self.brain_api_url}/simulations', json=simulation_data, ) print(f"📤 模拟提交状态: {sim_resp.status_code}") if sim_resp.status_code != 201: print(f"❌ {name} 提交失败: {sim_resp.text}") return {'name': name, 'success': False, 'error': f'提交失败: {sim_resp.status_code}'} # 获取进度URL并轮询结果 sim_progress_url = sim_resp.headers['location'] print(f"⏳ {name} 进度URL获取成功") while True: sim_progress_resp = self.client.get(sim_progress_url) retry_after_sec = float(sim_progress_resp.headers.get("Retry-After", 0)) if retry_after_sec == 0: # simulation done! break print(f"⏰ {name} 等待 {retry_after_sec} 秒...") sleep(retry_after_sec) # 获取最终的alpha ID result_data = sim_progress_resp.json() alpha_id = result_data["alpha"] print(f"✅ {name} 模拟完成! Alpha ID: {alpha_id}") return { 'name': name, 'expression': expression, 'alpha_id': alpha_id, 'success': True, 'result_data': result_data } except Exception as e: print(f"❌ {name} 模拟过程中出错: {str(e)}") return {'name': name, 'success': False, 'error': str(e)} def _print_summary(self, results): """打印测试结果汇总""" print("\n" + "=" * 50) print("📊 Alpha测试结果汇总") print("=" * 50) success_count = sum(1 for r in results if r.get('success', False)) failed_count = len(results) - success_count print(f"✅ 成功: {success_count} 个") print(f"❌ 失败: {failed_count} 个") print("\n成功详情:") for result in results: if result.get('success'): print(f" ✓ {result['name']}: {result['alpha_id']}") if failed_count > 0: print("\n失败详情:") for result in results: if not result.get('success'): print(f" ✗ {result['name']}: {result.get('error', '未知错误')}") def close(self): """关闭客户端连接""" if self.client: self.client.close() self.client = None # 使用示例 if __name__ == "__main__": # 创建API实例 wq_api = WorldQuantBrainAPI() try: # 1. 批量测试所有Alpha print("🎯 开始批量测试Alpha因子...") start_time = datetime.now() results = wq_api.batch_simulate_alphas( alphas_config_file='alphas.json', max_workers=2 # 并发数量，根据API限制调整 ) end_time = datetime.now() print(f"\n⏱️ 测试总耗时: {(end_time - start_time).total_seconds():.2f} 秒") # 2. 提取成功的Alpha ID successful_alphas = [r for r in results if r.get('success')] alpha_ids = [r['alpha_id'] for r in successful_alphas] print(f"\n📋 成功模拟的Alpha数量: {len(alpha_ids)}") for alpha in successful_alphas: print(f" - {alpha['name']}: {alpha['alpha_id']}") except Exception as e: print(f"❌ 程序执行出错: {str(e)}") import traceback traceback.print_exc() # 打印完整的错误堆栈 finally: # 确保连接关闭 wq_api.close()