244 lines
8.2 KiB
Python
244 lines
8.2 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""
|
|||
|
|
申万行业热度数据回溯补录工具
|
|||
|
|
读取本地CSV文件，存入MySQL sw_heat_daily表
|
|||
|
|
用于历史数据补录或重新导入
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import pandas as pd
|
|||
|
|
import pymysql
|
|||
|
|
import os
|
|||
|
|
import re
|
|||
|
|
from datetime import datetime
|
|||
|
|
import argparse
|
|||
|
|
|
|||
|
|
|
|||
|
|
def extract_date_from_filename(filename):
|
|||
|
|
"""
|
|||
|
|
从文件名提取日期
|
|||
|
|
例如: heat_ranking_2026-03-04_5d.csv -> 2026-03-04 -> 20260304
|
|||
|
|
"""
|
|||
|
|
# 匹配日期格式 YYYY-MM-DD
|
|||
|
|
pattern = r'(\d{4}-\d{2}-\d{2})'
|
|||
|
|
match = re.search(pattern, filename)
|
|||
|
|
|
|||
|
|
if match:
|
|||
|
|
date_str = match.group(1) # 2026-03-04
|
|||
|
|
# 转换为 YYYYMMDD 格式
|
|||
|
|
calc_date = date_str.replace('-', '') # 20260304
|
|||
|
|
return calc_date, date_str
|
|||
|
|
else:
|
|||
|
|
raise ValueError(f"无法从文件名提取日期: {filename}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def read_heat_csv(file_path):
|
|||
|
|
"""
|
|||
|
|
读取热度排名CSV文件
|
|||
|
|
"""
|
|||
|
|
if not os.path.exists(file_path):
|
|||
|
|
raise FileNotFoundError(f"文件不存在: {file_path}")
|
|||
|
|
|
|||
|
|
print(f"📂 读取文件: {file_path}")
|
|||
|
|
df = pd.read_csv(file_path, encoding='utf-8-sig')
|
|||
|
|
|
|||
|
|
# 检查必要列是否存在
|
|||
|
|
required_cols = ['code', 'name', 'total_return', 'avg_turnover',
|
|||
|
|
'volume_ratio', 'momentum', 'total_return_score',
|
|||
|
|
'avg_turnover_score', 'volume_ratio_score',
|
|||
|
|
'momentum_score', 'heat_score']
|
|||
|
|
|
|||
|
|
missing_cols = [col for col in required_cols if col not in df.columns]
|
|||
|
|
if missing_cols:
|
|||
|
|
# 尝试兼容旧版列名（如果有的话）
|
|||
|
|
col_mapping = {
|
|||
|
|
'行业代码': 'code',
|
|||
|
|
'行业名称': 'name',
|
|||
|
|
# 添加其他可能的映射
|
|||
|
|
}
|
|||
|
|
for old_col, new_col in col_mapping.items():
|
|||
|
|
if old_col in df.columns and new_col not in df.columns:
|
|||
|
|
df[new_col] = df[old_col]
|
|||
|
|
|
|||
|
|
# 再次检查
|
|||
|
|
still_missing = [col for col in required_cols if col not in df.columns]
|
|||
|
|
if still_missing:
|
|||
|
|
raise ValueError(f"CSV文件缺少必要列: {still_missing}")
|
|||
|
|
|
|||
|
|
print(f"✅ 成功读取 {len(df)} 条行业数据")
|
|||
|
|
return df
|
|||
|
|
|
|||
|
|
|
|||
|
|
def save_to_mysql(df, calc_date, db_config, dry_run=False):
|
|||
|
|
"""
|
|||
|
|
将DataFrame存入MySQL
|
|||
|
|
:param calc_date: YYYYMMDD 格式的日期字符串
|
|||
|
|
:param dry_run: 如果为True，只打印SQL不执行
|
|||
|
|
"""
|
|||
|
|
if dry_run:
|
|||
|
|
print(f"\n🔍 [DRY RUN] 预览模式，不实际写入数据库")
|
|||
|
|
print(f"📅 目标日期: {calc_date}")
|
|||
|
|
print(f"📊 数据预览:")
|
|||
|
|
print(df[['code', 'name', 'heat_score']].head())
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
conn = None
|
|||
|
|
try:
|
|||
|
|
conn = pymysql.connect(**db_config)
|
|||
|
|
cursor = conn.cursor()
|
|||
|
|
|
|||
|
|
# 检查该日期是否已有数据
|
|||
|
|
cursor.execute(
|
|||
|
|
"SELECT COUNT(*) FROM sw_heat_daily WHERE calc_date = %s",
|
|||
|
|
(calc_date,)
|
|||
|
|
)
|
|||
|
|
existing = cursor.fetchone()[0]
|
|||
|
|
|
|||
|
|
if existing > 0:
|
|||
|
|
print(f"⚠️ 日期 {calc_date} 已有 {existing} 条记录，将执行覆盖更新")
|
|||
|
|
|
|||
|
|
# 准备插入SQL
|
|||
|
|
insert_sql = """
|
|||
|
|
INSERT INTO sw_heat_daily
|
|||
|
|
(calc_date, code, name, total_return, avg_turnover, volume_ratio, momentum,
|
|||
|
|
total_return_score, avg_turnover_score, volume_ratio_score, momentum_score, heat_score)
|
|||
|
|
VALUES
|
|||
|
|
(%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|||
|
|
ON DUPLICATE KEY UPDATE
|
|||
|
|
name = VALUES(name),
|
|||
|
|
total_return = VALUES(total_return),
|
|||
|
|
avg_turnover = VALUES(avg_turnover),
|
|||
|
|
volume_ratio = VALUES(volume_ratio),
|
|||
|
|
momentum = VALUES(momentum),
|
|||
|
|
total_return_score = VALUES(total_return_score),
|
|||
|
|
avg_turnover_score = VALUES(avg_turnover_score),
|
|||
|
|
volume_ratio_score = VALUES(volume_ratio_score),
|
|||
|
|
momentum_score = VALUES(momentum_score),
|
|||
|
|
heat_score = VALUES(heat_score),
|
|||
|
|
created_at = CURRENT_TIMESTAMP;
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
# 转换数据为列表元组
|
|||
|
|
data_tuples = []
|
|||
|
|
for _, row in df.iterrows():
|
|||
|
|
data_tuples.append((
|
|||
|
|
calc_date,
|
|||
|
|
str(row['code']),
|
|||
|
|
str(row['name']),
|
|||
|
|
float(row['total_return']),
|
|||
|
|
float(row['avg_turnover']),
|
|||
|
|
float(row['volume_ratio']),
|
|||
|
|
float(row['momentum']),
|
|||
|
|
float(row['total_return_score']),
|
|||
|
|
float(row['avg_turnover_score']),
|
|||
|
|
float(row['volume_ratio_score']),
|
|||
|
|
float(row['momentum_score']),
|
|||
|
|
float(row['heat_score'])
|
|||
|
|
))
|
|||
|
|
|
|||
|
|
# 执行批量插入
|
|||
|
|
cursor.executemany(insert_sql, data_tuples)
|
|||
|
|
conn.commit()
|
|||
|
|
|
|||
|
|
action = "覆盖更新" if existing > 0 else "新增"
|
|||
|
|
print(f"✅ MySQL{action}成功: {calc_date} 共 {len(data_tuples)} 条记录")
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"❌ MySQL操作失败: {e}")
|
|||
|
|
if conn:
|
|||
|
|
conn.rollback()
|
|||
|
|
return False
|
|||
|
|
finally:
|
|||
|
|
if conn:
|
|||
|
|
conn.close()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def batch_import(data_dir, db_config, pattern=None, dry_run=False):
|
|||
|
|
"""
|
|||
|
|
批量导入目录下所有符合格式的CSV文件
|
|||
|
|
"""
|
|||
|
|
if pattern is None:
|
|||
|
|
pattern = r'heat_ranking_\d{4}-\d{2}-\d{2}_.*\.csv'
|
|||
|
|
|
|||
|
|
files = [f for f in os.listdir(data_dir) if re.match(pattern, f)]
|
|||
|
|
files.sort()
|
|||
|
|
|
|||
|
|
print(f"📁 发现 {len(files)} 个待导入文件")
|
|||
|
|
|
|||
|
|
success_count = 0
|
|||
|
|
for filename in files:
|
|||
|
|
file_path = os.path.join(data_dir, filename)
|
|||
|
|
try:
|
|||
|
|
calc_date, _ = extract_date_from_filename(filename)
|
|||
|
|
df = read_heat_csv(file_path)
|
|||
|
|
|
|||
|
|
if save_to_mysql(df, calc_date, db_config, dry_run):
|
|||
|
|
success_count += 1
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"❌ 处理 {filename} 失败: {e}")
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
print(f"\n📊 导入完成: {success_count}/{len(files)} 个文件成功")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
# 数据库配置（请根据实际情况修改）
|
|||
|
|
config = {
|
|||
|
|
'host': '10.127.2.207',
|
|||
|
|
'user': 'financial_prod',
|
|||
|
|
'password': 'mmTFncqmDal5HLRGY0BV',
|
|||
|
|
'database': 'reference',
|
|||
|
|
'port': 3306,
|
|||
|
|
'charset': 'utf8mb4'
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 命令行参数解析
|
|||
|
|
parser = argparse.ArgumentParser(description='申万行业热度数据回溯补录工具')
|
|||
|
|
parser.add_argument('file', nargs='?', help='单个CSV文件路径')
|
|||
|
|
parser.add_argument('--date', '-d', help='指定日期(YYYYMMDD)，覆盖文件名中的日期')
|
|||
|
|
parser.add_argument('--batch', '-b', action='store_true', help='批量导入目录下所有文件')
|
|||
|
|
parser.add_argument('--dir', default='./sw_data', help='数据目录路径(默认: ./sw_data)')
|
|||
|
|
parser.add_argument('--dry-run', action='store_true', help='预览模式，不实际写入数据库')
|
|||
|
|
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
# 批量模式
|
|||
|
|
if args.batch:
|
|||
|
|
batch_import(args.dir, config, dry_run=args.dry_run)
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 单文件模式
|
|||
|
|
if args.file:
|
|||
|
|
file_path = args.file
|
|||
|
|
else:
|
|||
|
|
# 默认示例文件
|
|||
|
|
file_path = './sw_data/heat_ranking_2026-03-04_5d.csv'
|
|||
|
|
if not os.path.exists(file_path):
|
|||
|
|
print("请提供CSV文件路径，或使用 --batch 批量导入")
|
|||
|
|
parser.print_help()
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 读取CSV
|
|||
|
|
df = read_heat_csv(file_path)
|
|||
|
|
|
|||
|
|
# 确定日期
|
|||
|
|
if args.date:
|
|||
|
|
calc_date = args.date # 命令行指定的日期
|
|||
|
|
print(f"📅 使用指定日期: {calc_date}")
|
|||
|
|
else:
|
|||
|
|
calc_date, original_date = extract_date_from_filename(file_path)
|
|||
|
|
print(f"📅 从文件名提取日期: {original_date} -> {calc_date}")
|
|||
|
|
|
|||
|
|
# 存入MySQL
|
|||
|
|
save_to_mysql(df, calc_date, config, dry_run=args.dry_run)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"❌ 错误: {e}")
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|