""" ETF 全量历史数据本地缓存 ======================== 一次性下载全市场 ETF（含已退市）的基础信息和日线数据到本地， 供回测中按 ref_date 截取历史数据，消除前视偏差。 用法: # 首次下载（约 30-60 分钟，取决于 API 限流） python scripts/etf_data_cache.py # 增量更新（只下载缺失的新数据） python scripts/etf_data_cache.py --update """ import os import sys import time import logging from pathlib import Path from datetime import datetime import pandas as pd sys.path.insert(0, str(Path(__file__).parent.parent)) from dotenv import load_dotenv load_dotenv() import tushare as ts logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s') logger = logging.getLogger(__name__) # 缓存目录 CACHE_DIR = Path(__file__).parent.parent / 'data' / 'etf_cache' DAILY_DIR = CACHE_DIR / 'daily' BASIC_PATH = CACHE_DIR / 'fund_basic.csv' class ETFDataCache: """ETF 全量历史数据缓存管理器""" def __init__(self): self.pro = ts.pro_api(os.getenv('TUSHARE_TOKEN')) CACHE_DIR.mkdir(parents=True, exist_ok=True) DAILY_DIR.mkdir(parents=True, exist_ok=True) self._basic_df = None # 懒加载 # ---------------------------------------------------------- # API 调用（带重试 + 限流） # ---------------------------------------------------------- def _api_call(self, func, **kwargs): for attempt in range(3): try: result = func(**kwargs) time.sleep(0.35) return result except Exception as e: if attempt < 2: wait = 2 * (attempt + 1) logger.warning(f" API 重试 ({attempt+1}/3): {e}, 等待 {wait}s") time.sleep(wait) else: raise # ---------------------------------------------------------- # 1. 下载并缓存 fund_basic # ---------------------------------------------------------- def download_basic(self, force: bool = False): """下载全量 ETF 基础信息（含已退市）""" if BASIC_PATH.exists() and not force: logger.info(f"fund_basic 缓存已存在: {BASIC_PATH}") return logger.info("下载全量 ETF 基础信息...") fields = 'ts_code,name,management,list_date,delist_date,fund_type,invest_type,benchmark,type,trustee,status' dfs = [] for status in ['L', 'D']: # L=上市, D=已退市 df = self._api_call(self.pro.fund_basic, market='E', status=status, fields=fields) if df is not None and not df.empty: dfs.append(df) logger.info(f" status={status}: {len(df)} 只") if not dfs: raise RuntimeError("获取 ETF 列表失败") basic = pd.concat(dfs, ignore_index=True).drop_duplicates(subset='ts_code') basic.to_csv(BASIC_PATH, index=False, encoding='utf-8-sig') logger.info(f"fund_basic 已保存: {len(basic)} 只 -> {BASIC_PATH}") # ---------------------------------------------------------- # 2. 批量下载日线数据 # ---------------------------------------------------------- def download_daily(self, force: bool = False): """批量下载所有 ETF 的全历史日线数据""" basic = self.load_basic() codes = basic['ts_code'].tolist() total = len(codes) logger.info(f"准备下载 {total} 只 ETF 的日线数据...") downloaded = 0 skipped = 0 failed = 0 for i, code in enumerate(codes): csv_path = DAILY_DIR / f"{code}.csv" if csv_path.exists() and not force: # 增量更新: 读取已有数据的最后日期 try: existing = pd.read_csv(csv_path, nrows=1) # 只读首行检查 if not existing.empty: skipped += 1 continue except Exception: pass if (i - skipped) % 20 == 0: logger.info(f" 进度: {i}/{total} (下载={downloaded}, 跳过={skipped}, 失败={failed})") try: df = self._api_call( self.pro.fund_daily, ts_code=code, fields='ts_code,trade_date,open,high,low,close,vol,amount' ) if df is not None and not df.empty: df = df.sort_values('trade_date') df.to_csv(csv_path, index=False) downloaded += 1 else: failed += 1 except Exception as e: logger.warning(f" {code} 下载失败: {e}") failed += 1 logger.info(f"日线数据下载完成: 下载={downloaded}, 跳过={skipped}, 失败={failed}") def update_daily(self): """增量更新: 只为已有缓存文件追加新数据""" basic = self.load_basic() codes = basic['ts_code'].tolist() today_str = datetime.now().strftime('%Y%m%d') updated = 0 for code in codes: csv_path = DAILY_DIR / f"{code}.csv" if not csv_path.exists(): continue try: existing = pd.read_csv(csv_path) if existing.empty: continue last_date = str(existing['trade_date'].max()) if last_date >= today_str: continue # 下载 last_date 之后的数据 new_df = self._api_call( self.pro.fund_daily, ts_code=code, start_date=str(int(last_date) + 1), end_date=today_str, fields='ts_code,trade_date,open,high,low,close,vol,amount' ) if new_df is not None and not new_df.empty: combined = pd.concat([existing, new_df], ignore_index=True) combined = combined.drop_duplicates(subset='trade_date').sort_values('trade_date') combined.to_csv(csv_path, index=False) updated += 1 except Exception: pass logger.info(f"增量更新完成: {updated} 只有新数据") # ---------------------------------------------------------- # 3. 数据读取接口（回测用） # ---------------------------------------------------------- def load_basic(self) -> pd.DataFrame: """加载 fund_basic 缓存""" if self._basic_df is not None: return self._basic_df if not BASIC_PATH.exists(): raise FileNotFoundError(f"fund_basic 缓存不存在，请先运行: python scripts/etf_data_cache.py") self._basic_df = pd.read_csv(BASIC_PATH) return self._basic_df def load_cached_daily(self, ts_code: str, end_date: str = None) -> pd.DataFrame: """ 加载某只 ETF 的日线数据，截至 end_date（含）。 Args: ts_code: ETF 代码 end_date: 截止日期 YYYYMMDD，None 表示全部 Returns: DataFrame with columns [trade_date, open, high, low, close, vol, amount] 按 trade_date 升序排列 """ csv_path = DAILY_DIR / f"{ts_code}.csv" if not csv_path.exists(): return pd.DataFrame() df = pd.read_csv(csv_path) if df.empty: return df df['trade_date'] = df['trade_date'].astype(str) df = df.sort_values('trade_date') if end_date: end_str = str(end_date).replace('-', '') df = df[df['trade_date'] <= end_str] return df def load_cached_daily_as_series(self, ts_code: str, end_date: str = None, column: str = 'close') -> pd.Series: """加载某只 ETF 的单列数据，index 为 datetime""" df = self.load_cached_daily(ts_code, end_date) if df.empty: return pd.Series(dtype=float) df['date'] = pd.to_datetime(df['trade_date']) return df.set_index('date')[column].astype(float) def load_cached_ohlcv(self, ts_code: str, end_date: str = None) -> pd.DataFrame: """加载 OHLCV 数据，index 为 datetime（与 动量.py 的 all_data 格式兼容）""" df = self.load_cached_daily(ts_code, end_date) if df.empty: return pd.DataFrame() df['date'] = pd.to_datetime(df['trade_date']) df = df.set_index('date').sort_index() df = df.rename(columns={'vol': 'volume'}) return df[['open', 'high', 'low', 'close', 'volume']].astype(float) def ensure_downloaded(self): """确保基础信息和日线数据都已下载""" self.download_basic() self.download_daily() def get_available_codes_at(self, ref_date: str) -> list: """获取在 ref_date 时已上市且未退市的 ETF 代码列表""" basic = self.load_basic() basic['list_date'] = basic['list_date'].astype(str) mask = basic['list_date'] <= ref_date # 排除在 ref_date 之前已退市的 if 'delist_date' in basic.columns: delist = basic['delist_date'].astype(str).fillna('99991231') mask = mask & (delist > ref_date) return basic[mask]['ts_code'].tolist() # ---------------------------------------------------------- # CLI # ---------------------------------------------------------- if __name__ == '__main__': import argparse parser = argparse.ArgumentParser(description='ETF 全量历史数据缓存下载') parser.add_argument('--update', action='store_true', help='增量更新已有缓存') parser.add_argument('--force', action='store_true', help='强制重新下载全部') args = parser.parse_args() cache = ETFDataCache() if args.update: cache.download_basic(force=True) cache.update_daily() else: cache.download_basic(force=args.force) cache.download_daily(force=args.force) # 统计 basic = cache.load_basic() n_daily = len(list(DAILY_DIR.glob('*.csv'))) logger.info(f"\n缓存统计: fund_basic={len(basic)} 只, 日线文件={n_daily} 个")