276 lines
8.4 KiB
Python
276 lines
8.4 KiB
Python
"""
|
||
A股历史数据获取系统
|
||
功能：获取所有A股从2010年至今的历史行情数据
|
||
"""
|
||
|
||
import tushare as ts
|
||
import pandas as pd
|
||
import os
|
||
import time
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
|
||
# 配置
|
||
BASE_DIR = Path(__file__).parent
|
||
DATA_DIR = BASE_DIR / 'data'
|
||
LOGS_DIR = BASE_DIR / 'logs'
|
||
STOCK_LIST_FILE = BASE_DIR / 'A股股票列表.csv'
|
||
|
||
# 数据保存目录
|
||
DATA_DIR.mkdir(exist_ok=True)
|
||
LOGS_DIR.mkdir(exist_ok=True)
|
||
|
||
# 时间范围
|
||
START_DATE = '20100101'
|
||
END_DATE = datetime.now().strftime('%Y%m%d')
|
||
|
||
# 请求间隔（秒）- tushare积分限制
|
||
REQUEST_INTERVAL = 5
|
||
|
||
|
||
def setup_tushare(token=None):
|
||
"""初始化tushare"""
|
||
import os
|
||
|
||
# 优先级：参数 > 环境变量 > 配置文件
|
||
if not token:
|
||
token = os.environ.get('TUSHARE_TOKEN', '')
|
||
|
||
if not token:
|
||
# 尝试从配置文件读取
|
||
config_file = BASE_DIR / 'config.txt'
|
||
if config_file.exists():
|
||
token = config_file.read_text().strip()
|
||
|
||
if not token:
|
||
print("错误：未设置 Tushare Token！")
|
||
print("请通过以下方式之一设置：")
|
||
print(" 1. 设置环境变量: export TUSHARE_TOKEN=your_token")
|
||
print(" 2. 创建配置文件: echo 'your_token' > config.txt")
|
||
print(" 3. 注册地址: https://tushare.pro/register")
|
||
raise ValueError("缺少 Tushare Token")
|
||
|
||
ts.set_token(token)
|
||
return ts.pro_api()
|
||
|
||
|
||
def load_stock_list():
|
||
"""加载股票列表"""
|
||
df = pd.read_csv(STOCK_LIST_FILE)
|
||
# 清理列名中的空白
|
||
df.columns = df.columns.str.strip()
|
||
print(f"加载股票列表: {len(df)} 只股票")
|
||
return df
|
||
|
||
|
||
def get_stock_codes_with_suffix(df):
|
||
"""将股票代码转换为tushare格式（添加后缀）
|
||
|
||
市场代码规则：
|
||
- 6开头 → SH（上海）
|
||
- 0、3开头 → SZ（深圳）
|
||
- 4、8开头 → BJ（北京）
|
||
"""
|
||
codes = []
|
||
for code in df['code']:
|
||
code = str(code).zfill(6) # 补零到6位
|
||
first_digit = code[0]
|
||
|
||
if first_digit == '6':
|
||
ts_code = f"{code}.SH"
|
||
elif first_digit in ('0', '3'):
|
||
ts_code = f"{code}.SZ"
|
||
elif first_digit in ('4', '8'):
|
||
ts_code = f"{code}.BJ"
|
||
else:
|
||
# 未知市场，默认深圳
|
||
ts_code = f"{code}.SZ"
|
||
|
||
codes.append(ts_code)
|
||
return codes
|
||
|
||
|
||
def fetch_daily_data(pro, codes, start_date, end_date):
|
||
"""逐个获取日线数据（每次一支股票），支持断点续传"""
|
||
total = len(codes)
|
||
|
||
# 加载已完成的股票列表
|
||
completed_file = DATA_DIR / 'completed_stocks.txt'
|
||
completed_stocks = set()
|
||
if completed_file.exists():
|
||
lines = completed_file.read_text().strip().split('\n')
|
||
completed_stocks = set(line.strip() for line in lines if line.strip())
|
||
print(f"已完成: {len(completed_stocks)} 只股票")
|
||
|
||
# 统计已有数据
|
||
existing_data_file = DATA_DIR / 'stock_daily_data.parquet'
|
||
if existing_data_file.exists():
|
||
existing_df = pd.read_parquet(existing_data_file)
|
||
print(f"已有数据: {len(existing_df)} 条记录")
|
||
|
||
print(f"\n共 {total} 只股票，待处理: {total - len(completed_stocks)} 只")
|
||
print(f"预计耗时: {(total - len(completed_stocks)) * REQUEST_INTERVAL / 60:.1f} 分钟")
|
||
print("-" * 50)
|
||
|
||
for i, ts_code in enumerate(codes):
|
||
# 跳过已完成的
|
||
if ts_code in completed_stocks:
|
||
print(f"[{i+1}/{total}] {ts_code} 已完成，跳过")
|
||
continue
|
||
|
||
try:
|
||
print(f"[{i+1}/{total}] 获取 {ts_code}...", end=' ', flush=True)
|
||
|
||
df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
|
||
|
||
if df is not None and len(df) > 0:
|
||
print(f"成功，{len(df)} 条记录")
|
||
# 实时保存单只股票数据
|
||
save_single_stock(df, ts_code, completed_file)
|
||
else:
|
||
print("无数据")
|
||
# 无数据也标记为完成
|
||
mark_completed(ts_code, completed_file)
|
||
|
||
except Exception as e:
|
||
print(f"错误: {e}")
|
||
# 出错不标记完成，下次重试
|
||
|
||
# 每次请求后休息
|
||
if i < total - 1:
|
||
time.sleep(REQUEST_INTERVAL)
|
||
|
||
# 最后合并所有数据
|
||
return merge_all_data()
|
||
|
||
|
||
def save_single_stock(df, ts_code, completed_file):
|
||
"""保存单只股票数据并标记完成"""
|
||
# 读取已有数据
|
||
output_file = DATA_DIR / 'stock_daily_data.parquet'
|
||
|
||
if output_file.exists():
|
||
existing_df = pd.read_parquet(output_file)
|
||
# 删除该股票的旧数据（如果有）
|
||
existing_df = existing_df[existing_df['ts_code'] != ts_code]
|
||
# 合并新数据
|
||
combined_df = pd.concat([existing_df, df], ignore_index=True)
|
||
else:
|
||
combined_df = df
|
||
|
||
# 排序
|
||
combined_df = combined_df.sort_values(['ts_code', 'trade_date']).reset_index(drop=True)
|
||
|
||
# 保存
|
||
combined_df.to_parquet(output_file, index=False)
|
||
|
||
# 标记完成
|
||
mark_completed(ts_code, completed_file)
|
||
|
||
|
||
def merge_all_data():
|
||
"""最后合并所有数据（用于返回）"""
|
||
output_file = DATA_DIR / 'stock_daily_data.parquet'
|
||
if output_file.exists():
|
||
return [pd.read_parquet(output_file)]
|
||
return []
|
||
|
||
|
||
def save_progress(all_data, ts_code, completed_file):
|
||
"""实时保存进度（保留兼容性）"""
|
||
# 合并并保存数据
|
||
combined_df = pd.concat(all_data, ignore_index=True)
|
||
combined_df = combined_df.sort_values(['ts_code', 'trade_date']).reset_index(drop=True)
|
||
|
||
# 保存parquet
|
||
output_file = DATA_DIR / 'stock_daily_data.parquet'
|
||
combined_df.to_parquet(output_file, index=False)
|
||
|
||
# 标记完成
|
||
mark_completed(ts_code, completed_file)
|
||
|
||
|
||
def mark_completed(ts_code, completed_file):
|
||
"""标记股票已完成"""
|
||
with open(completed_file, 'a') as f:
|
||
f.write(ts_code + '\n')
|
||
|
||
|
||
def save_to_parquet(df, filename):
|
||
"""保存为parquet格式（高效压缩）"""
|
||
filepath = DATA_DIR / filename
|
||
df.to_parquet(filepath, index=False)
|
||
print(f"保存到: {filepath}")
|
||
print(f"文件大小: {filepath.stat().st_size / 1024 / 1024:.2f} MB")
|
||
|
||
|
||
def save_to_csv(df, filename):
|
||
"""保存为CSV格式"""
|
||
filepath = DATA_DIR / filename
|
||
df.to_csv(filepath, index=False)
|
||
print(f"保存到: {filepath}")
|
||
print(f"文件大小: {filepath.stat().st_size / 1024 / 1024:.2f} MB")
|
||
|
||
|
||
def main():
|
||
"""主函数"""
|
||
print("=" * 60)
|
||
print("A股历史数据获取系统")
|
||
print("=" * 60)
|
||
print(f"数据时间范围: {START_DATE} ~ {END_DATE}")
|
||
print(f"数据保存目录: {DATA_DIR}")
|
||
print("=" * 60)
|
||
|
||
# 初始化tushare
|
||
print("\n初始化 Tushare...")
|
||
pro = setup_tushare()
|
||
|
||
# 加载股票列表
|
||
print("\n加载股票列表...")
|
||
stock_df = load_stock_list()
|
||
codes = get_stock_codes_with_suffix(stock_df)
|
||
print(f"共 {len(codes)} 只股票")
|
||
|
||
# 获取日线数据
|
||
print("\n开始获取日线数据...")
|
||
all_data = fetch_daily_data(pro, codes, START_DATE, END_DATE)
|
||
|
||
if all_data:
|
||
# 合并所有数据
|
||
print("\n合并数据...")
|
||
combined_df = pd.concat(all_data, ignore_index=True)
|
||
print(f"总记录数: {len(combined_df)}")
|
||
|
||
# 按日期排序
|
||
combined_df = combined_df.sort_values(['ts_code', 'trade_date']).reset_index(drop=True)
|
||
|
||
# 保存数据
|
||
print("\n保存数据...")
|
||
timestamp = datetime.now().strftime('%Y%m%d')
|
||
|
||
# 保存为parquet（推荐，压缩率高）
|
||
save_to_parquet(combined_df, f'A股日线数据_{timestamp}.parquet')
|
||
|
||
# 同时保存为CSV（方便查看）
|
||
save_to_csv(combined_df, f'A股日线数据_{timestamp}.csv')
|
||
|
||
# 显示数据概览
|
||
print("\n数据概览:")
|
||
print(f" 股票数量: {combined_df['ts_code'].nunique()}")
|
||
print(f" 日期范围: {combined_df['trade_date'].min()} ~ {combined_df['trade_date'].max()}")
|
||
print(f" 总记录数: {len(combined_df)}")
|
||
print("\n列名:")
|
||
print(combined_df.columns.tolist())
|
||
print("\n前5条数据:")
|
||
print(combined_df.head())
|
||
else:
|
||
print("未获取到任何数据")
|
||
|
||
print("\n" + "=" * 60)
|
||
print("数据获取完成！")
|
||
print("=" * 60)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|