""" 完整数据更新脚本 同时更新K线数据和财务数据，支持分批处理和进度显示 """ import sys import os import asyncio import logging from datetime import date, datetime sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from src.storage.database import db_manager from src.storage.stock_repository import StockRepository from src.data.data_manager import DataManager from src.config.settings import Settings # 配置日志 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) def convert_to_baostock_format(stock_code: str) -> str: """ 将6位股票代码转换为Baostock格式（9位） Args: stock_code: 6位股票代码 Returns: 9位Baostock格式股票代码 """ if len(stock_code) == 6: # 判断市场类型 if stock_code.startswith(('6', '9')): return f"sh.{stock_code}" elif stock_code.startswith(('0', '3')): return f"sz.{stock_code}" else: return stock_code return stock_code async def update_kline_data_batch(stocks: list, data_manager: DataManager, repository: StockRepository, batch_size: int = 10): """ 分批更新K线数据 Args: stocks: 股票列表 data_manager: 数据管理器 repository: 存储库 batch_size: 每批处理的股票数量 Returns: 更新结果 """ total_kline_data = [] success_count = 0 error_count = 0 # 分批处理 for i in range(0, len(stocks), batch_size): batch = stocks[i:i + batch_size] logger.info(f"处理K线数据批次 {i//batch_size + 1}/{(len(stocks)-1)//batch_size + 1}: {len(batch)}只股票") batch_kline_data = [] batch_success = 0 batch_error = 0 for stock in batch: try: # 转换为Baostock格式 baostock_code = convert_to_baostock_format(stock.code) # 获取K线数据（最近3个月） end_date = date.today() start_date = date(end_date.year, end_date.month - 3, 1) kline_data = await data_manager.get_daily_kline_data( baostock_code, start_date, end_date ) if kline_data: # 将数据中的代码转换回6位格式 for data in kline_data: data["code"] = stock.code batch_kline_data.extend(kline_data) batch_success += 1 logger.info(f"股票{stock.code}获取到{len(kline_data)}条K线数据") else: logger.warning(f"股票{stock.code}未获取到K线数据") batch_error += 1 except Exception as e: logger.error(f"获取股票{stock.code}K线数据失败: {str(e)}") batch_error += 1 continue # 小延迟避免请求过快 await asyncio.sleep(0.2) # 保存当前批次的数据 if batch_kline_data: try: save_result = repository.save_daily_kline_data(batch_kline_data) logger.info(f"批次K线数据保存结果: {save_result}") total_kline_data.extend(batch_kline_data) except Exception as e: logger.error(f"保存批次K线数据失败: {str(e)}") batch_error += len(batch) success_count += batch_success error_count += batch_error logger.info(f"批次完成: 成功{batch_success}只, 失败{batch_error}只") return { "success": True, "total_stocks": len(stocks), "success_count": success_count, "error_count": error_count, "kline_data_count": len(total_kline_data) } async def update_financial_data_batch(stocks: list, data_manager: DataManager, repository: StockRepository, batch_size: int = 10): """ 分批更新财务数据 Args: stocks: 股票列表 data_manager: 数据管理器 repository: 存储库 batch_size: 每批处理的股票数量 Returns: 更新结果 """ total_financial_data = [] success_count = 0 error_count = 0 # 设置测试年份和季度 test_year = 2023 test_quarter = 4 # 分批处理 for i in range(0, len(stocks), batch_size): batch = stocks[i:i + batch_size] logger.info(f"处理财务数据批次 {i//batch_size + 1}/{(len(stocks)-1)//batch_size + 1}: {len(batch)}只股票") batch_financial_data = [] batch_success = 0 batch_error = 0 for stock in batch: try: # 转换为Baostock格式 baostock_code = convert_to_baostock_format(stock.code) # 获取财务数据 financial_data = await data_manager.get_financial_report( baostock_code, test_year, test_quarter ) if financial_data: # 将数据中的代码转换回6位格式 for data in financial_data: data["code"] = stock.code batch_financial_data.extend(financial_data) batch_success += 1 logger.info(f"股票{stock.code}获取到{len(financial_data)}条财务数据") else: logger.warning(f"股票{stock.code}未获取到财务数据") batch_error += 1 except Exception as e: logger.error(f"获取股票{stock.code}财务数据失败: {str(e)}") batch_error += 1 continue # 小延迟避免请求过快 await asyncio.sleep(0.3) # 保存当前批次的数据 if batch_financial_data: try: save_result = repository.save_financial_report_data(batch_financial_data) logger.info(f"批次财务数据保存结果: {save_result}") total_financial_data.extend(batch_financial_data) except Exception as e: logger.error(f"保存批次财务数据失败: {str(e)}") batch_error += len(batch) success_count += batch_success error_count += batch_error logger.info(f"批次完成: 成功{batch_success}只, 失败{batch_error}只") return { "success": True, "total_stocks": len(stocks), "success_count": success_count, "error_count": error_count, "financial_data_count": len(total_financial_data) } async def update_all_data(): """ 更新所有数据（K线数据和财务数据） """ try: logger.info("开始更新所有股票数据...") # 加载配置 settings = Settings() logger.info("配置加载成功") # 创建数据管理器 data_manager = DataManager() logger.info("数据管理器创建成功") # 创建存储库 repository = StockRepository(db_manager.get_session()) logger.info("存储库创建成功") # 获取股票基础信息 stocks = repository.get_stock_basic_info() logger.info(f"获取到{len(stocks)}只股票基础信息") if not stocks: logger.error("没有股票基础信息，无法更新数据") return {"success": False, "error": "没有股票基础信息"} # 选择前50只股票进行测试（避免处理时间过长） test_stocks = stocks[:50] test_codes = [stock.code for stock in test_stocks] logger.info(f"测试股票代码: {test_codes}") # 更新K线数据 logger.info("=== 开始更新K线数据 ===") kline_result = await update_kline_data_batch(test_stocks, data_manager, repository, batch_size=5) # 更新财务数据 logger.info("=== 开始更新财务数据 ===") financial_result = await update_financial_data_batch(test_stocks, data_manager, repository, batch_size=5) # 汇总结果 result = { "success": True, "kline_data": kline_result, "financial_data": financial_result, "total_stocks": len(test_stocks) } logger.info(f"所有数据更新完成: {result}") return result except Exception as e: logger.error(f"数据更新异常: {str(e)}") return {"success": False, "error": str(e)} def main(): """ 主函数 """ logger.info("开始完整数据更新流程...") # 运行异步更新 result = asyncio.run(update_all_data()) if result["success"]: logger.info("数据更新成功！") kline_result = result["kline_data"] financial_result = result["financial_data"] print("=== 数据更新结果汇总 ===") print(f"处理股票总数: {result['total_stocks']}") print("\n=== K线数据更新结果 ===") print(f"✓ 成功股票数: {kline_result['success_count']}") print(f"✓ 失败股票数: {kline_result['error_count']}") print(f"✓ 获取K线数据: {kline_result['kline_data_count']}条") print("\n=== 财务数据更新结果 ===") print(f"✓ 成功股票数: {financial_result['success_count']}") print(f"✓ 失败股票数: {financial_result['error_count']}") print(f"✓ 获取财务数据: {financial_result['financial_data_count']}条") print("\n=== 数据库验证 ===") # 验证数据库中的数据 try: repository = StockRepository(db_manager.get_session()) # 查询K线数据 kline_count = repository.session.query(repository.DailyKline).count() print(f"✓ 日K线数据表: {kline_count}条记录") # 查询财务数据 financial_count = repository.session.query(repository.FinancialReport).count() print(f"✓ 财务报告表: {financial_count}条记录") # 查询股票基础信息 stock_count = repository.session.query(repository.StockBasicInfo).count() print(f"✓ 股票基础信息: {stock_count}条记录") except Exception as e: print(f"⚠ 数据库验证失败: {str(e)}") print("\n数据更新流程完成！") else: logger.error("数据更新失败！") print(f"更新失败: {result.get('error')}") return result if __name__ == "__main__": main()