🛠️ feat(data_downloader): 添加交易日历表创建、数据下载、保存与更新功能
➕ feat(sql): 新增交易日历表的SQL创建语句
This commit is contained in:
parent
541d39139b
commit
77d990794b
@ -1,3 +1,5 @@
|
||||
import argparse
|
||||
|
||||
import pandas as pd
|
||||
import tushare as ts
|
||||
import numpy as np
|
||||
@ -871,23 +873,154 @@ def perform_incremental_update(processes=4):
|
||||
print(f"总计: {len(stocks_to_update)} 只股票")
|
||||
|
||||
|
||||
def main():
|
||||
"""主函数，允许用户选择更新方式"""
|
||||
import argparse
|
||||
# 添加创建日历表的函数
|
||||
def create_calendar_table(engine):
|
||||
"""创建交易日历表"""
|
||||
# 从sql文件读取创建表的SQL
|
||||
with open('sql/calendar.sql', 'r', encoding='utf-8') as f:
|
||||
create_table_sql = f.read()
|
||||
|
||||
# 命令行参数解析
|
||||
parser = argparse.ArgumentParser(description='股票数据更新工具')
|
||||
parser.add_argument('--init', action='store_true', help='初始化元数据表')
|
||||
parser.add_argument('--mode', choices=['full', 'incremental', 'both'],
|
||||
default='full', help='更新模式: full=全量更新, incremental=增量更新, both=两者都执行')
|
||||
parser.add_argument('--year', type=int, default=2020, help='全量更新的起始年份')
|
||||
parser.add_argument('--processes', type=int, default=8, help='使用的进程数量')
|
||||
parser.add_argument('--resume', action='store_true', help='仅更新全量更新失败的股票,即status=2部分')
|
||||
with engine.connect() as conn:
|
||||
conn.execute(text(create_table_sql))
|
||||
conn.commit()
|
||||
print("交易日历表创建成功")
|
||||
|
||||
|
||||
# 添加下载日历数据的函数
|
||||
def download_calendar_data(pro, start_date=None, end_date=None):
|
||||
"""
|
||||
下载交易日历数据
|
||||
|
||||
Args:
|
||||
pro: Tushare API对象
|
||||
start_date: 开始日期，格式YYYYMMDD，默认为2000年初
|
||||
end_date: 结束日期，格式YYYYMMDD，默认为当前日期之后的一年
|
||||
|
||||
Returns:
|
||||
pandas.DataFrame: 交易日历数据
|
||||
"""
|
||||
# 默认获取从2000年至今后一年的日历数据
|
||||
if start_date is None:
|
||||
start_date = '20000101'
|
||||
|
||||
if end_date is None:
|
||||
# 获取当前日期后一年的数据，以便预先有足够的交易日历
|
||||
current_date = datetime.now()
|
||||
next_year = current_date + timedelta(days=365)
|
||||
end_date = next_year.strftime('%Y%m%d')
|
||||
|
||||
print(f"正在下载从 {start_date} 到 {end_date} 的交易日历数据...")
|
||||
|
||||
try:
|
||||
# 调用Tushare的trade_cal接口获取交易日历
|
||||
df = pro.trade_cal(exchange='SSE', start_date=start_date, end_date=end_date)
|
||||
|
||||
# 重命名列以符合数据库表结构
|
||||
df = df.rename(columns={
|
||||
'cal_date': 'trade_date',
|
||||
'exchange': 'exchange',
|
||||
'is_open': 'is_open',
|
||||
'pretrade_date': 'pretrade_date'
|
||||
})
|
||||
|
||||
# 将日期格式从YYYYMMDD转换为YYYY-MM-DD
|
||||
df['trade_date'] = pd.to_datetime(df['trade_date']).dt.strftime('%Y-%m-%d')
|
||||
if 'pretrade_date' in df.columns:
|
||||
df['pretrade_date'] = pd.to_datetime(df['pretrade_date']).dt.strftime('%Y-%m-%d')
|
||||
|
||||
# 添加创建和更新时间
|
||||
current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
||||
df['create_time'] = current_time
|
||||
df['update_time'] = current_time
|
||||
|
||||
print(f"成功下载 {len(df)} 条交易日历记录")
|
||||
return df
|
||||
|
||||
except Exception as e:
|
||||
print(f"下载交易日历数据时出错: {e}")
|
||||
return pd.DataFrame()
|
||||
|
||||
|
||||
# 添加保存日历数据的函数
|
||||
def save_calendar_data(calendar_df, engine):
|
||||
"""
|
||||
将交易日历数据保存到数据库
|
||||
|
||||
Args:
|
||||
calendar_df: 包含交易日历数据的DataFrame
|
||||
engine: SQLAlchemy引擎
|
||||
"""
|
||||
if calendar_df.empty:
|
||||
print("没有交易日历数据可保存")
|
||||
return
|
||||
|
||||
try:
|
||||
# 使用to_sql将数据写入数据库，如果表存在则替换
|
||||
calendar_df.to_sql('calendar', engine, if_exists='replace', index=False)
|
||||
print(f"成功保存 {len(calendar_df)} 条交易日历数据到数据库")
|
||||
except Exception as e:
|
||||
print(f"保存交易日历数据时出错: {e}")
|
||||
|
||||
|
||||
# 添加交易日历更新函数
|
||||
def update_calendar(engine, pro):
|
||||
"""
|
||||
更新交易日历数据
|
||||
|
||||
Args:
|
||||
engine: SQLAlchemy引擎
|
||||
pro: Tushare API对象
|
||||
"""
|
||||
try:
|
||||
# 检查calendar表是否存在
|
||||
with engine.connect() as conn:
|
||||
result = conn.execute(text("SHOW TABLES LIKE 'calendar'"))
|
||||
table_exists = result.fetchone() is not None
|
||||
|
||||
if not table_exists:
|
||||
# 如果表不存在，创建表并下载所有历史数据
|
||||
create_calendar_table(engine)
|
||||
calendar_df = download_calendar_data(pro)
|
||||
save_calendar_data(calendar_df, engine)
|
||||
else:
|
||||
# 如果表存在，查询最新的交易日期
|
||||
with engine.connect() as conn:
|
||||
result = conn.execute(text("SELECT MAX(trade_date) FROM calendar"))
|
||||
latest_date = result.fetchone()[0]
|
||||
|
||||
if latest_date:
|
||||
# 将日期转换为YYYYMMDD格式
|
||||
start_date = (datetime.strptime(str(latest_date), '%Y-%m-%d') + timedelta(days=1)).strftime('%Y%m%d')
|
||||
# 下载新数据
|
||||
calendar_df = download_calendar_data(pro, start_date=start_date)
|
||||
|
||||
if not calendar_df.empty:
|
||||
# 将新数据追加到表中
|
||||
calendar_df.to_sql('calendar', engine, if_exists='append', index=False)
|
||||
print(f"成功更新 {len(calendar_df)} 条新的交易日历数据")
|
||||
else:
|
||||
print("没有新的交易日历数据需要更新")
|
||||
else:
|
||||
# 如果表为空，下载所有历史数据
|
||||
calendar_df = download_calendar_data(pro)
|
||||
save_calendar_data(calendar_df, engine)
|
||||
|
||||
except Exception as e:
|
||||
print(f"更新交易日历时出错: {e}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='股票数据下载器')
|
||||
parser.add_argument('--init', action='store_true', help='初始化元数据')
|
||||
parser.add_argument('--mode', choices=['full', 'incremental', 'both'], default='full',
|
||||
help='更新模式: full(全量), incremental(增量), both(两者都)')
|
||||
parser.add_argument('--year', type=int, default=2020, help='起始年份(用于全量更新)')
|
||||
parser.add_argument('--processes', type=int, default=4, help='并行进程数')
|
||||
parser.add_argument('--resume', action='store_true', help='是否从中断处继续')
|
||||
parser.add_argument('--update-calendar', action='store_true', help='更新交易日历数据')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# 如果需要初始化元数据
|
||||
if args.init:
|
||||
# 加载配置
|
||||
config = load_config()
|
||||
# 创建数据库引擎
|
||||
@ -896,6 +1029,8 @@ def main():
|
||||
ts.set_token(config['tushare_token'])
|
||||
pro = ts.pro_api()
|
||||
|
||||
# 如果需要初始化元数据
|
||||
if args.init:
|
||||
# 创建元数据表
|
||||
create_metadata_table(engine)
|
||||
# 下载并保存股票元数据
|
||||
@ -903,6 +1038,15 @@ def main():
|
||||
save_metadata(metadata_df, engine)
|
||||
print("元数据初始化完成")
|
||||
|
||||
# 初始化交易日历数据
|
||||
update_calendar(engine, pro)
|
||||
print("交易日历初始化完成")
|
||||
|
||||
# 如果需要更新交易日历
|
||||
if args.update_calendar:
|
||||
update_calendar(engine, pro)
|
||||
print("交易日历更新完成")
|
||||
|
||||
# 根据选择的模式执行更新
|
||||
if args.mode == 'full' or args.mode == 'both':
|
||||
perform_full_update(args.year, args.processes, args.resume)
|
||||
|
||||
12
sql/calendar.sql
Normal file
12
sql/calendar.sql
Normal file
@ -0,0 +1,12 @@
|
||||
CREATE TABLE IF NOT EXISTS `calendar` (
|
||||
trade_date DATE PRIMARY KEY COMMENT '交易日期',
|
||||
exchange VARCHAR(10) COMMENT '交易所代码',
|
||||
is_open TINYINT COMMENT '是否交易（0：休市 1：交易）',
|
||||
pretrade_date DATE COMMENT '上一个交易日',
|
||||
|
||||
create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
|
||||
update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
|
||||
|
||||
INDEX idx_exchange (exchange),
|
||||
INDEX idx_is_open (is_open)
|
||||
) ENGINE=InnoDB COMMENT '交易日历表';
|
||||
Loading…
Reference in New Issue
Block a user