feat: 添加热点股追踪数据导出/导入脚本
This commit is contained in:
Sakurasan
Executable
+213
@@ -0,0 +1,213 @@
|
||||
#!/usr/bin/env python3
|
||||
"""热点股追踪数据导出/导入脚本
|
||||
|
||||
用法：
|
||||
导出8月20日数据：
|
||||
python scripts/export_hotspot_data.py export --date 2026-08-20
|
||||
|
||||
导入数据：
|
||||
python scripts/export_hotspot_data.py import --file hotspot_2026-08-20.json
|
||||
|
||||
查看数据库中有哪些日期的数据：
|
||||
python scripts/export_hotspot_data.py list
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
# 添加项目路径
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "backend"))
|
||||
|
||||
from database import get_connection, dict_from_row
|
||||
|
||||
|
||||
def export_data(date: str, output_file: Optional[str] = None) -> dict:
|
||||
"""导出指定日期的热点股追踪数据"""
|
||||
conn = get_connection()
|
||||
try:
|
||||
# 1. 导出每日题材前20
|
||||
themes = conn.execute(
|
||||
"SELECT * FROM daily_top_themes WHERE trade_date = ? ORDER BY rank ASC",
|
||||
(date,),
|
||||
).fetchall()
|
||||
|
||||
# 2. 导出每日核心股票
|
||||
core_stocks = conn.execute(
|
||||
"SELECT * FROM daily_core_stocks WHERE trade_date = ? ORDER BY rank ASC",
|
||||
(date,),
|
||||
).fetchall()
|
||||
|
||||
# 3. 导出每日核心股票与题材的关联
|
||||
stock_themes = conn.execute(
|
||||
"SELECT * FROM daily_core_stock_themes WHERE trade_date = ?",
|
||||
(date,),
|
||||
).fetchall()
|
||||
|
||||
data = {
|
||||
"version": "1.0",
|
||||
"export_time": datetime.now().isoformat(),
|
||||
"trade_date": date,
|
||||
"themes": [dict_from_row(r) for r in themes],
|
||||
"core_stocks": [dict_from_row(r) for r in core_stocks],
|
||||
"stock_themes": [dict_from_row(r) for r in stock_themes],
|
||||
}
|
||||
|
||||
# 统计信息
|
||||
data["stats"] = {
|
||||
"theme_count": len(data["themes"]),
|
||||
"core_stock_count": len(data["core_stocks"]),
|
||||
"stock_theme_count": len(data["stock_themes"]),
|
||||
}
|
||||
|
||||
# 保存到文件
|
||||
if output_file is None:
|
||||
output_file = f"hotspot_{date}.json"
|
||||
|
||||
with open(output_file, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||||
|
||||
print(f"✅ 导出成功：{output_file}")
|
||||
print(f" 题材数量：{data['stats']['theme_count']}")
|
||||
print(f" 核心股票：{data['stats']['core_stock_count']}")
|
||||
print(f" 股票-题材关联：{data['stats']['stock_theme_count']}")
|
||||
|
||||
return data
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def import_data(input_file: str, dry_run: bool = False) -> bool:
|
||||
"""导入热点股追踪数据"""
|
||||
if not os.path.exists(input_file):
|
||||
print(f"❌ 文件不存在：{input_file}")
|
||||
return False
|
||||
|
||||
with open(input_file, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
|
||||
# 验证数据格式
|
||||
required_keys = ["version", "trade_date", "themes", "core_stocks", "stock_themes"]
|
||||
if not all(k in data for k in required_keys):
|
||||
print("❌ 数据格式错误：缺少必要字段")
|
||||
return False
|
||||
|
||||
date = data["trade_date"]
|
||||
print(f"📅 导入日期：{date}")
|
||||
print(f" 题材数量：{len(data['themes'])}")
|
||||
print(f" 核心股票：{len(data['core_stocks'])}")
|
||||
print(f" 股票-题材关联：{len(data['stock_themes'])}")
|
||||
|
||||
if dry_run:
|
||||
print("\n🔍 试运行模式，不写入数据库")
|
||||
return True
|
||||
|
||||
# 检查是否已存在
|
||||
conn = get_connection()
|
||||
try:
|
||||
existing = conn.execute(
|
||||
"SELECT 1 FROM daily_core_stocks WHERE trade_date = ? LIMIT 1",
|
||||
(date,),
|
||||
).fetchone()
|
||||
|
||||
if existing:
|
||||
print(f"\n⚠️ {date} 的数据已存在，将跳过重复数据（使用 INSERT OR IGNORE）")
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
# 写入数据库
|
||||
conn = get_connection()
|
||||
try:
|
||||
# 导入每日题材前20
|
||||
for item in data["themes"]:
|
||||
# 移除自动生成的id和created_at
|
||||
item.pop("id", None)
|
||||
item.pop("created_at", None)
|
||||
conn.execute(
|
||||
"INSERT OR IGNORE INTO daily_top_themes (trade_date, theme_code, theme_name, bf3, hot_rank, rank) VALUES (?,?,?,?,?,?)",
|
||||
(item["trade_date"], item["theme_code"], item["theme_name"], item["bf3"], item["hot_rank"], item["rank"]),
|
||||
)
|
||||
|
||||
# 导入每日核心股票
|
||||
for item in data["core_stocks"]:
|
||||
# 移除自动生成的id和created_at
|
||||
item.pop("id", None)
|
||||
item.pop("created_at", None)
|
||||
conn.execute(
|
||||
"INSERT OR IGNORE INTO daily_core_stocks (trade_date, stock_code, stock_name, f3, cover_count, rank) VALUES (?,?,?,?,?,?)",
|
||||
(item["trade_date"], item["stock_code"], item["stock_name"], item["f3"], item["cover_count"], item["rank"]),
|
||||
)
|
||||
|
||||
# 导入每日核心股票与题材的关联
|
||||
for item in data["stock_themes"]:
|
||||
# 移除自动生成的id和created_at
|
||||
item.pop("id", None)
|
||||
item.pop("created_at", None)
|
||||
conn.execute(
|
||||
"INSERT OR IGNORE INTO daily_core_stock_themes (trade_date, stock_code, theme_code, theme_name) VALUES (?,?,?,?)",
|
||||
(item["trade_date"], item["stock_code"], item["theme_code"], item["theme_name"]),
|
||||
)
|
||||
|
||||
conn.commit()
|
||||
print(f"\n✅ 导入成功：{date}")
|
||||
return True
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
print(f"\n❌ 导入失败：{e}")
|
||||
return False
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def list_dates() -> list[str]:
|
||||
"""列出数据库中所有有数据的日期"""
|
||||
conn = get_connection()
|
||||
try:
|
||||
rows = conn.execute(
|
||||
"SELECT DISTINCT trade_date FROM daily_core_stocks ORDER BY trade_date DESC"
|
||||
).fetchall()
|
||||
return [row["trade_date"] for row in rows]
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="热点股追踪数据导出/导入工具")
|
||||
subparsers = parser.add_subparsers(dest="command", help="可用命令")
|
||||
|
||||
# 导出命令
|
||||
export_parser = subparsers.add_parser("export", help="导出数据")
|
||||
export_parser.add_argument("--date", required=True, help="交易日期 (YYYY-MM-DD)")
|
||||
export_parser.add_argument("--output", "-o", help="输出文件路径")
|
||||
|
||||
# 导入命令
|
||||
import_parser = subparsers.add_parser("import", help="导入数据")
|
||||
import_parser.add_argument("--file", "-f", required=True, help="输入文件路径")
|
||||
import_parser.add_argument("--dry-run", action="store_true", help="试运行，不写入数据库")
|
||||
|
||||
# 列表命令
|
||||
subparsers.add_parser("list", help="列出所有有数据的日期")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.command == "export":
|
||||
export_data(args.date, args.output)
|
||||
elif args.command == "import":
|
||||
import_data(args.file, args.dry_run)
|
||||
elif args.command == "list":
|
||||
dates = list_dates()
|
||||
if dates:
|
||||
print("📅 数据库中的日期：")
|
||||
for d in dates:
|
||||
print(f" {d}")
|
||||
else:
|
||||
print("📭 数据库中暂无数据")
|
||||
else:
|
||||
parser.print_help()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user