809 lines
29 KiB
Python
809 lines
29 KiB
Python
import pandas as pd
|
||
from sqlalchemy import text
|
||
import argparse
|
||
import tushare as ts
|
||
from datetime import date
|
||
from collections import deque
|
||
import time
|
||
|
||
from stock_config import TUSHARE_TOKEN, get_engine
|
||
|
||
# --- 数据库配置 ---
|
||
# 请根据实际情况修改数据库连接信息
|
||
# DB_CONN_STR = 'mysql+pymysql://root:password@127.0.0.1:3306/stock_data'
|
||
LOCAL_START_DATE = date(2018, 1, 1)
|
||
TUSHARE_LIMIT_LIST_START_DATE = date(2020, 1, 1)
|
||
TUSHARE_LIMIT_LIST_MAX_CALLS_PER_MINUTE = 200
|
||
COMBINED_TABLE = "stock_limit_combined_data"
|
||
|
||
|
||
class MinuteRateLimiter:
|
||
def __init__(self, max_calls):
|
||
self.max_calls = max_calls
|
||
self.call_times = deque()
|
||
|
||
def wait(self):
|
||
now = time.monotonic()
|
||
while self.call_times and now - self.call_times[0] >= 60:
|
||
self.call_times.popleft()
|
||
if len(self.call_times) >= self.max_calls:
|
||
sleep_seconds = 60 - (now - self.call_times[0]) + 0.1
|
||
print(f"⏳ Tushare 限流等待 {sleep_seconds:.1f} 秒，避免超过每分钟 {self.max_calls} 次。")
|
||
time.sleep(sleep_seconds)
|
||
now = time.monotonic()
|
||
while self.call_times and now - self.call_times[0] >= 60:
|
||
self.call_times.popleft()
|
||
self.call_times.append(time.monotonic())
|
||
|
||
|
||
def normalize_trade_dates(trade_dates):
|
||
return sorted({pd.to_datetime(trade_date).date() for trade_date in trade_dates})
|
||
|
||
|
||
def filter_dates_by_start(trade_dates, start_date, source_name):
|
||
normalized_dates = normalize_trade_dates(trade_dates)
|
||
filtered_dates = [trade_date for trade_date in normalized_dates if trade_date >= start_date]
|
||
skipped_count = len(normalized_dates) - len(filtered_dates)
|
||
if skipped_count:
|
||
print(f"ℹ️ {source_name} 跳过 {skipped_count} 个早于 {start_date} 的交易日。")
|
||
return filtered_dates
|
||
|
||
|
||
def init_db(engine, full_maintenance=False):
|
||
"""
|
||
初始化数据库表结构
|
||
"""
|
||
with engine.begin() as conn:
|
||
print(f" - 确保 {COMBINED_TABLE} 存在 ...")
|
||
conn.execute(text(get_combined_table_sql()))
|
||
optimize_combined_table(conn, full_maintenance=full_maintenance)
|
||
|
||
|
||
def get_combined_table_sql():
|
||
return f"""
|
||
CREATE TABLE IF NOT EXISTS `{COMBINED_TABLE}` (
|
||
`id` BIGINT AUTO_INCREMENT PRIMARY KEY,
|
||
`trade_date` DATE NOT NULL COMMENT '交易日期',
|
||
`symbol` VARCHAR(10) NOT NULL COMMENT '股票代码，以stock_daily为准',
|
||
`ts_code` VARCHAR(20) NOT NULL COMMENT 'TS股票代码，以stock_daily为准',
|
||
`name` VARCHAR(50) NOT NULL COMMENT '股票名称，以stock_daily关联基础表为准',
|
||
`final_status` TINYINT NOT NULL COMMENT '融合状态：1-涨停，2-炸板',
|
||
`source_count` INT NOT NULL DEFAULT 0 COMMENT '命中的数据源数量',
|
||
`source_flags` VARCHAR(50) COMMENT '命中的数据源：calc,tushare',
|
||
`calc_status` TINYINT COMMENT '本地行情计算状态：1-涨停，2-炸板',
|
||
`calc_limit_price` DECIMAL(16, 4) COMMENT '本地行情计算涨停价',
|
||
`calc_high` DECIMAL(16, 4) COMMENT 'stock_daily最高价',
|
||
`calc_close` DECIMAL(16, 4) COMMENT 'stock_daily收盘价',
|
||
`calc_pre_close` DECIMAL(16, 4) COMMENT 'stock_daily前收盘价',
|
||
`calc_vol` BIGINT COMMENT 'stock_daily成交量',
|
||
`market` VARCHAR(20) COMMENT '所属板块',
|
||
`tushare_limit_type` VARCHAR(1) COMMENT 'Tushare-D跌停U涨停Z炸板',
|
||
`tushare_industry` VARCHAR(100) COMMENT 'Tushare-所属行业',
|
||
`tushare_name` VARCHAR(50) COMMENT 'Tushare-名称',
|
||
`tushare_close` DECIMAL(16, 4) COMMENT 'Tushare-收盘价',
|
||
`tushare_pct_chg` DECIMAL(10, 4) COMMENT 'Tushare-涨跌幅',
|
||
`tushare_amount` DECIMAL(24, 4) COMMENT 'Tushare-成交额',
|
||
`tushare_limit_amount` DECIMAL(24, 4) COMMENT 'Tushare-板上成交金额',
|
||
`tushare_float_mv` DECIMAL(24, 4) COMMENT 'Tushare-流通市值',
|
||
`tushare_total_mv` DECIMAL(24, 4) COMMENT 'Tushare-总市值',
|
||
`tushare_turnover_ratio` DECIMAL(10, 4) COMMENT 'Tushare-换手率',
|
||
`tushare_fd_amount` DECIMAL(24, 4) COMMENT 'Tushare-封单金额',
|
||
`tushare_first_time` VARCHAR(20) COMMENT 'Tushare-首次封板时间',
|
||
`tushare_last_time` VARCHAR(20) COMMENT 'Tushare-最后封板时间',
|
||
`tushare_open_times` INT COMMENT 'Tushare-炸板/开板次数',
|
||
`tushare_up_stat` VARCHAR(50) COMMENT 'Tushare-涨停统计',
|
||
`tushare_limit_times` INT COMMENT 'Tushare-连板数',
|
||
UNIQUE KEY `uniq_date_code` (`trade_date`, `ts_code`),
|
||
KEY `idx_trade_date` (`trade_date`),
|
||
KEY `idx_symbol` (`symbol`)
|
||
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;
|
||
"""
|
||
|
||
|
||
def get_table_columns(conn, table_name):
|
||
return {
|
||
row[0]
|
||
for row in conn.execute(text(f"SHOW COLUMNS FROM `{table_name}`")).fetchall()
|
||
}
|
||
|
||
|
||
def get_column_collation(conn, table_name, column_name):
|
||
row = conn.execute(
|
||
text(
|
||
"""
|
||
SELECT collation_name
|
||
FROM information_schema.columns
|
||
WHERE table_schema = DATABASE()
|
||
AND table_name = :table_name
|
||
AND column_name = :column_name
|
||
"""
|
||
),
|
||
{"table_name": table_name, "column_name": column_name},
|
||
).first()
|
||
return row[0] if row else None
|
||
|
||
|
||
def index_exists(conn, table_name, index_name):
|
||
return conn.execute(
|
||
text(
|
||
"""
|
||
SELECT 1
|
||
FROM information_schema.statistics
|
||
WHERE table_schema = DATABASE()
|
||
AND table_name = :table_name
|
||
AND index_name = :index_name
|
||
LIMIT 1
|
||
"""
|
||
),
|
||
{"table_name": table_name, "index_name": index_name},
|
||
).first() is not None
|
||
|
||
|
||
def get_index_definition(conn, table_name, index_name):
|
||
rows = conn.execute(
|
||
text(
|
||
"""
|
||
SELECT column_name, non_unique
|
||
FROM information_schema.statistics
|
||
WHERE table_schema = DATABASE()
|
||
AND table_name = :table_name
|
||
AND index_name = :index_name
|
||
ORDER BY seq_in_index
|
||
"""
|
||
),
|
||
{"table_name": table_name, "index_name": index_name},
|
||
).fetchall()
|
||
if not rows:
|
||
return None
|
||
return {
|
||
"columns": [row[0] for row in rows],
|
||
"unique": rows[0][1] == 0,
|
||
}
|
||
|
||
|
||
def index_matches(conn, table_name, index_name, columns, unique=False):
|
||
index_def = get_index_definition(conn, table_name, index_name)
|
||
if not index_def:
|
||
return False
|
||
return index_def["columns"] == columns and (not unique or index_def["unique"])
|
||
|
||
|
||
def ensure_combined_columns(conn):
|
||
existing_columns = get_table_columns(conn, COMBINED_TABLE)
|
||
columns = {
|
||
"calc_vol": "BIGINT COMMENT 'stock_daily成交量'",
|
||
}
|
||
|
||
for column_name, column_def in columns.items():
|
||
if column_name not in existing_columns:
|
||
conn.execute(
|
||
text(
|
||
f"ALTER TABLE `{COMBINED_TABLE}` "
|
||
f"ADD COLUMN `{column_name}` {column_def} AFTER `calc_pre_close`"
|
||
)
|
||
)
|
||
|
||
ensure_combined_join_collation(conn)
|
||
|
||
|
||
def ensure_combined_join_collation(conn):
|
||
daily_collation = get_column_collation(conn, "stock_daily", "ts_code")
|
||
combined_collation = get_column_collation(conn, COMBINED_TABLE, "ts_code")
|
||
if not daily_collation or not combined_collation or daily_collation == combined_collation:
|
||
return
|
||
|
||
print(f" - 对齐 ts_code collation：{combined_collation} -> {daily_collation} ...")
|
||
conn.execute(
|
||
text(
|
||
f"""
|
||
ALTER TABLE `{COMBINED_TABLE}`
|
||
MODIFY `ts_code` VARCHAR(20)
|
||
CHARACTER SET utf8mb4 COLLATE {daily_collation}
|
||
NOT NULL COMMENT 'TS股票代码，以stock_daily为准'
|
||
"""
|
||
)
|
||
)
|
||
|
||
|
||
def sync_combined_vol_from_daily(conn):
|
||
conn.execute(
|
||
text(
|
||
f"""
|
||
UPDATE `{COMBINED_TABLE}` c
|
||
JOIN `stock_daily` d
|
||
ON c.ts_code = d.ts_code
|
||
AND c.trade_date = d.trade_date
|
||
SET c.symbol = d.symbol,
|
||
c.calc_vol = d.vol
|
||
"""
|
||
)
|
||
)
|
||
|
||
|
||
def dedupe_combined_table(conn):
|
||
conn.execute(
|
||
text(
|
||
f"""
|
||
DELETE c
|
||
FROM `{COMBINED_TABLE}` c
|
||
LEFT JOIN `stock_daily` d
|
||
ON c.ts_code = d.ts_code
|
||
AND c.trade_date = d.trade_date
|
||
WHERE d.ts_code IS NULL
|
||
"""
|
||
)
|
||
)
|
||
conn.execute(
|
||
text(
|
||
f"""
|
||
DELETE c1
|
||
FROM `{COMBINED_TABLE}` c1
|
||
JOIN `{COMBINED_TABLE}` c2
|
||
ON c1.trade_date = c2.trade_date
|
||
AND c1.ts_code = c2.ts_code
|
||
WHERE
|
||
(
|
||
COALESCE(c1.source_count, 0) * 10
|
||
+ IF(c1.calc_status IS NOT NULL, 4, 0)
|
||
+ IF(c1.tushare_limit_type IS NOT NULL, 2, 0)
|
||
+ IF(c1.calc_vol IS NOT NULL, 1, 0)
|
||
) <
|
||
(
|
||
COALESCE(c2.source_count, 0) * 10
|
||
+ IF(c2.calc_status IS NOT NULL, 4, 0)
|
||
+ IF(c2.tushare_limit_type IS NOT NULL, 2, 0)
|
||
+ IF(c2.calc_vol IS NOT NULL, 1, 0)
|
||
)
|
||
OR (
|
||
(
|
||
COALESCE(c1.source_count, 0) * 10
|
||
+ IF(c1.calc_status IS NOT NULL, 4, 0)
|
||
+ IF(c1.tushare_limit_type IS NOT NULL, 2, 0)
|
||
+ IF(c1.calc_vol IS NOT NULL, 1, 0)
|
||
) =
|
||
(
|
||
COALESCE(c2.source_count, 0) * 10
|
||
+ IF(c2.calc_status IS NOT NULL, 4, 0)
|
||
+ IF(c2.tushare_limit_type IS NOT NULL, 2, 0)
|
||
+ IF(c2.calc_vol IS NOT NULL, 1, 0)
|
||
)
|
||
AND c1.id < c2.id
|
||
)
|
||
"""
|
||
)
|
||
)
|
||
|
||
|
||
def ensure_combined_indexes(conn, full_maintenance=False):
|
||
needs_unique_rebuild = index_exists(
|
||
conn, COMBINED_TABLE, "uniq_date_code"
|
||
) and not index_matches(
|
||
conn, COMBINED_TABLE, "uniq_date_code", ["trade_date", "ts_code"], unique=True
|
||
)
|
||
needs_unique_create = not index_matches(
|
||
conn, COMBINED_TABLE, "uniq_date_code", ["trade_date", "ts_code"], unique=True
|
||
)
|
||
|
||
if full_maintenance or needs_unique_rebuild or needs_unique_create:
|
||
print(" - 清理历史重复/无效融合数据 ...")
|
||
dedupe_combined_table(conn)
|
||
|
||
if index_exists(conn, COMBINED_TABLE, "uniq_date_code") and not index_matches(
|
||
conn, COMBINED_TABLE, "uniq_date_code", ["trade_date", "ts_code"], unique=True
|
||
):
|
||
print(" - 重建 uniq_date_code 唯一索引 ...")
|
||
conn.execute(text(f"ALTER TABLE `{COMBINED_TABLE}` DROP INDEX `uniq_date_code`"))
|
||
|
||
if not index_matches(
|
||
conn, COMBINED_TABLE, "uniq_date_code", ["trade_date", "ts_code"], unique=True
|
||
):
|
||
print(" - 创建 uniq_date_code 唯一索引 ...")
|
||
conn.execute(
|
||
text(
|
||
f"ALTER TABLE `{COMBINED_TABLE}` "
|
||
"ADD UNIQUE KEY `uniq_date_code` (`trade_date`, `ts_code`)"
|
||
)
|
||
)
|
||
if not index_exists(conn, COMBINED_TABLE, "idx_trade_date"):
|
||
print(" - 创建 idx_trade_date 索引 ...")
|
||
conn.execute(
|
||
text(f"ALTER TABLE `{COMBINED_TABLE}` ADD KEY `idx_trade_date` (`trade_date`)")
|
||
)
|
||
if not index_exists(conn, COMBINED_TABLE, "idx_symbol"):
|
||
print(" - 创建 idx_symbol 索引 ...")
|
||
conn.execute(
|
||
text(f"ALTER TABLE `{COMBINED_TABLE}` ADD KEY `idx_symbol` (`symbol`)")
|
||
)
|
||
if not index_exists(conn, COMBINED_TABLE, "idx_trade_date_status"):
|
||
print(" - 创建 idx_trade_date_status 索引 ...")
|
||
conn.execute(
|
||
text(
|
||
f"ALTER TABLE `{COMBINED_TABLE}` "
|
||
"ADD KEY `idx_trade_date_status` (`trade_date`, `final_status`)"
|
||
)
|
||
)
|
||
|
||
|
||
def optimize_combined_table(conn, full_maintenance=False):
|
||
print(" - 检查字段 ...")
|
||
ensure_combined_columns(conn)
|
||
if full_maintenance:
|
||
print(" - 从 stock_daily 全量回填 calc_vol ...")
|
||
sync_combined_vol_from_daily(conn)
|
||
else:
|
||
print(" - 跳过全量 calc_vol 回填，当前批次写入时会直接带入 vol。")
|
||
print(" - 检查索引 ...")
|
||
ensure_combined_indexes(conn, full_maintenance=full_maintenance)
|
||
print(" - 校验重复键 ...")
|
||
validate_no_duplicate_combined_rows(conn)
|
||
|
||
|
||
def validate_no_duplicate_combined_rows(conn):
|
||
duplicate_count = conn.execute(
|
||
text(
|
||
f"""
|
||
SELECT COUNT(*)
|
||
FROM (
|
||
SELECT trade_date, ts_code
|
||
FROM `{COMBINED_TABLE}`
|
||
GROUP BY trade_date, ts_code
|
||
HAVING COUNT(*) > 1
|
||
) duplicated
|
||
"""
|
||
)
|
||
).scalar()
|
||
if duplicate_count:
|
||
raise RuntimeError(f"{COMBINED_TABLE} 存在 {duplicate_count} 组重复 trade_date+ts_code 数据")
|
||
|
||
|
||
def validate_no_duplicate_frame_rows(df, frame_name):
|
||
if df.empty:
|
||
return
|
||
missing_keys = df["trade_date"].isna() | df["ts_code"].isna()
|
||
if missing_keys.any():
|
||
examples = df.loc[missing_keys, ["trade_date", "ts_code"]].head(5).to_dict("records")
|
||
raise ValueError(f"{frame_name} 存在空 trade_date 或 ts_code，示例: {examples}")
|
||
duplicated = df.duplicated(subset=["trade_date", "ts_code"], keep=False)
|
||
if duplicated.any():
|
||
examples = df.loc[duplicated, ["trade_date", "ts_code"]].head(5).to_dict("records")
|
||
raise ValueError(f"{frame_name} 存在重复 trade_date+ts_code，示例: {examples}")
|
||
|
||
|
||
def fetch_tushare_limit_list(trade_dates):
|
||
tushare_trade_dates = filter_dates_by_start(
|
||
trade_dates,
|
||
TUSHARE_LIMIT_LIST_START_DATE,
|
||
"Tushare limit_list_d",
|
||
)
|
||
if not tushare_trade_dates:
|
||
print("ℹ️ 无需调用 Tushare limit_list_d。")
|
||
return pd.DataFrame(columns=["trade_date", "symbol", "ts_code"])
|
||
|
||
ts.set_token(TUSHARE_TOKEN)
|
||
pro = ts.pro_api()
|
||
|
||
frames = []
|
||
rate_limiter = MinuteRateLimiter(TUSHARE_LIMIT_LIST_MAX_CALLS_PER_MINUTE)
|
||
total_calls = len(tushare_trade_dates) * 2
|
||
current_call = 0
|
||
fields = (
|
||
"trade_date,ts_code,industry,name,close,pct_chg,amount,limit_amount,"
|
||
"float_mv,total_mv,turnover_ratio,fd_amount,first_time,last_time,"
|
||
"open_times,up_stat,limit_times,limit"
|
||
)
|
||
print(
|
||
f"🔎 开始获取 Tushare limit_list_d：{len(tushare_trade_dates)} 个交易日，"
|
||
f"预计 {total_calls} 次调用，限速 {TUSHARE_LIMIT_LIST_MAX_CALLS_PER_MINUTE} 次/分钟。"
|
||
)
|
||
for trade_date in tushare_trade_dates:
|
||
tushare_date = pd.to_datetime(trade_date).strftime("%Y%m%d")
|
||
for limit_type in ("U", "Z"):
|
||
current_call += 1
|
||
print(f" [{current_call}/{total_calls}] 获取 {tushare_date} limit_type={limit_type} ...")
|
||
rate_limiter.wait()
|
||
try:
|
||
ts_df = pro.limit_list_d(
|
||
trade_date=tushare_date,
|
||
limit_type=limit_type,
|
||
fields=fields,
|
||
)
|
||
except Exception as exc:
|
||
print(f"⚠️ Tushare limit_list_d 获取失败: {tushare_date} {limit_type}, 错误: {exc}")
|
||
continue
|
||
|
||
if ts_df is None or ts_df.empty:
|
||
print(f" [{current_call}/{total_calls}] {tushare_date} {limit_type} 无数据。")
|
||
continue
|
||
|
||
normalized_df = normalize_tushare_limit_list(ts_df)
|
||
if not normalized_df.empty:
|
||
print(f" [{current_call}/{total_calls}] {tushare_date} {limit_type} 获取 {len(normalized_df)} 条。")
|
||
frames.append(normalized_df)
|
||
|
||
if not frames:
|
||
print("ℹ️ Tushare limit_list_d 未返回有效数据。")
|
||
return pd.DataFrame(columns=["trade_date", "symbol", "ts_code"])
|
||
|
||
result_df = pd.concat(frames, ignore_index=True).drop_duplicates(
|
||
subset=["trade_date", "ts_code", "tushare_limit_type"],
|
||
keep="first",
|
||
)
|
||
print(f"✅ Tushare limit_list_d 获取完成：{len(result_df)} 条有效记录。")
|
||
return result_df
|
||
|
||
|
||
def normalize_tushare_limit_list(ts_df):
|
||
column_map = {
|
||
"limit": "tushare_limit_type",
|
||
"industry": "tushare_industry",
|
||
"name": "tushare_name",
|
||
"close": "tushare_close",
|
||
"pct_chg": "tushare_pct_chg",
|
||
"amount": "tushare_amount",
|
||
"limit_amount": "tushare_limit_amount",
|
||
"float_mv": "tushare_float_mv",
|
||
"total_mv": "tushare_total_mv",
|
||
"turnover_ratio": "tushare_turnover_ratio",
|
||
"fd_amount": "tushare_fd_amount",
|
||
"first_time": "tushare_first_time",
|
||
"last_time": "tushare_last_time",
|
||
"open_times": "tushare_open_times",
|
||
"up_stat": "tushare_up_stat",
|
||
"limit_times": "tushare_limit_times",
|
||
}
|
||
|
||
required_columns = {"trade_date", "ts_code", "limit"}
|
||
if not required_columns.issubset(ts_df.columns):
|
||
missing_columns = sorted(required_columns - set(ts_df.columns))
|
||
print(f"⚠️ Tushare 返回字段缺失，已跳过: {', '.join(missing_columns)}")
|
||
return pd.DataFrame(columns=["trade_date", "symbol", "ts_code"])
|
||
|
||
available_map = {
|
||
source: target for source, target in column_map.items() if source in ts_df.columns
|
||
}
|
||
normalized_df = ts_df[["trade_date", "ts_code", *available_map]].rename(
|
||
columns=available_map
|
||
).copy()
|
||
normalized_df["trade_date"] = pd.to_datetime(
|
||
normalized_df["trade_date"], format="%Y%m%d"
|
||
).dt.date
|
||
normalized_df["ts_code"] = normalized_df["ts_code"].astype(str).str.strip()
|
||
normalized_df["symbol"] = normalized_df["ts_code"].str.split(".").str[0].str.zfill(6)
|
||
|
||
for column in [
|
||
"tushare_close",
|
||
"tushare_pct_chg",
|
||
"tushare_amount",
|
||
"tushare_limit_amount",
|
||
"tushare_float_mv",
|
||
"tushare_total_mv",
|
||
"tushare_turnover_ratio",
|
||
"tushare_fd_amount",
|
||
"tushare_open_times",
|
||
"tushare_limit_times",
|
||
]:
|
||
if column in normalized_df.columns:
|
||
normalized_df[column] = pd.to_numeric(normalized_df[column], errors="coerce")
|
||
|
||
for column in [
|
||
"tushare_limit_type",
|
||
"tushare_industry",
|
||
"tushare_name",
|
||
"tushare_first_time",
|
||
"tushare_last_time",
|
||
"tushare_up_stat",
|
||
]:
|
||
if column in normalized_df.columns:
|
||
normalized_df[column] = normalized_df[column].astype(str)
|
||
|
||
return normalized_df
|
||
|
||
|
||
def build_combined_limit_data(daily_df, calc_df, tushare_df):
|
||
tushare_columns = [
|
||
"tushare_limit_type",
|
||
"tushare_industry",
|
||
"tushare_name",
|
||
"tushare_close",
|
||
"tushare_pct_chg",
|
||
"tushare_amount",
|
||
"tushare_limit_amount",
|
||
"tushare_float_mv",
|
||
"tushare_total_mv",
|
||
"tushare_turnover_ratio",
|
||
"tushare_fd_amount",
|
||
"tushare_first_time",
|
||
"tushare_last_time",
|
||
"tushare_open_times",
|
||
"tushare_up_stat",
|
||
"tushare_limit_times",
|
||
]
|
||
|
||
for column in tushare_columns:
|
||
if column not in tushare_df.columns:
|
||
tushare_df[column] = None
|
||
|
||
calc_key_df = pd.DataFrame(columns=["trade_date", "symbol"])
|
||
if not calc_df.empty:
|
||
calc_key_df = calc_df[["trade_date", "symbol"]].copy()
|
||
tushare_key_df = pd.DataFrame(columns=["trade_date", "symbol"])
|
||
if not tushare_df.empty:
|
||
tushare_key_df = tushare_df[["trade_date", "symbol"]].copy()
|
||
|
||
keys_df = pd.concat(
|
||
[calc_key_df, tushare_key_df],
|
||
ignore_index=True,
|
||
).drop_duplicates()
|
||
if keys_df.empty:
|
||
return pd.DataFrame()
|
||
|
||
for frame in [daily_df, calc_df, tushare_df, keys_df]:
|
||
if "trade_date" in frame.columns:
|
||
frame["trade_date"] = pd.to_datetime(frame["trade_date"]).dt.date
|
||
if "symbol" in frame.columns:
|
||
frame["symbol"] = frame["symbol"].astype(str).str.zfill(6)
|
||
keys_df = keys_df.drop_duplicates(subset=["trade_date", "symbol"])
|
||
|
||
canonical_df = daily_df[[
|
||
"trade_date",
|
||
"symbol",
|
||
"ts_code",
|
||
"name",
|
||
"high",
|
||
"close",
|
||
"pre_close",
|
||
"vol",
|
||
"market",
|
||
]].drop_duplicates(subset=["trade_date", "symbol"])
|
||
combined_df = keys_df.merge(
|
||
canonical_df,
|
||
how="inner",
|
||
on=["trade_date", "symbol"],
|
||
)
|
||
if combined_df.empty:
|
||
return pd.DataFrame()
|
||
|
||
if not calc_df.empty:
|
||
calc_merge_df = calc_df[[
|
||
"trade_date",
|
||
"symbol",
|
||
"status",
|
||
"limit_price",
|
||
]].rename(columns={
|
||
"status": "calc_status",
|
||
"limit_price": "calc_limit_price",
|
||
})
|
||
combined_df = combined_df.merge(
|
||
calc_merge_df,
|
||
how="left",
|
||
on=["trade_date", "symbol"],
|
||
)
|
||
else:
|
||
combined_df["calc_status"] = None
|
||
combined_df["calc_limit_price"] = None
|
||
|
||
if not tushare_df.empty:
|
||
tushare_merge_df = tushare_df[
|
||
["trade_date", "symbol", *tushare_columns]
|
||
].sort_values("tushare_limit_type").drop_duplicates(
|
||
subset=["trade_date", "symbol"],
|
||
keep="last",
|
||
)
|
||
combined_df = combined_df.merge(
|
||
tushare_merge_df,
|
||
how="left",
|
||
on=["trade_date", "symbol"],
|
||
)
|
||
else:
|
||
for column in tushare_columns:
|
||
combined_df[column] = None
|
||
|
||
combined_df["calc_high"] = combined_df["high"]
|
||
combined_df["calc_close"] = combined_df["close"]
|
||
combined_df["calc_pre_close"] = combined_df["pre_close"]
|
||
combined_df["calc_vol"] = combined_df["vol"]
|
||
|
||
combined_df["final_status"] = combined_df.apply(resolve_final_status, axis=1)
|
||
combined_df["source_flags"] = combined_df.apply(resolve_source_flags, axis=1)
|
||
combined_df["source_count"] = combined_df["source_flags"].apply(
|
||
lambda value: 0 if not value else len(value.split(","))
|
||
)
|
||
combined_df["name"] = combined_df["name"].astype(str).str.strip()
|
||
|
||
result_df = combined_df[[
|
||
"trade_date",
|
||
"symbol",
|
||
"ts_code",
|
||
"name",
|
||
"final_status",
|
||
"source_count",
|
||
"source_flags",
|
||
"calc_status",
|
||
"calc_limit_price",
|
||
"calc_high",
|
||
"calc_close",
|
||
"calc_pre_close",
|
||
"calc_vol",
|
||
"market",
|
||
*tushare_columns,
|
||
]].copy()
|
||
result_df = result_df.sort_values(
|
||
["trade_date", "ts_code", "source_count"],
|
||
ascending=[True, True, False],
|
||
)
|
||
return result_df.drop_duplicates(subset=["trade_date", "ts_code"], keep="first")
|
||
|
||
|
||
def resolve_final_status(row):
|
||
if pd.notna(row.get("calc_status")):
|
||
return int(row["calc_status"])
|
||
if row.get("tushare_limit_type") == "Z":
|
||
return 2
|
||
return 1
|
||
|
||
|
||
def resolve_source_flags(row):
|
||
flags = []
|
||
if pd.notna(row.get("calc_status")):
|
||
flags.append("calc")
|
||
if pd.notna(row.get("tushare_limit_type")):
|
||
flags.append("tushare")
|
||
return ",".join(flags)
|
||
|
||
|
||
def analyze_and_save_limits(start_date, end_date):
|
||
"""
|
||
计算并在数据库中落地涨停、炸板详情
|
||
"""
|
||
engine = get_engine()
|
||
|
||
# 0. 初始化检查表结构
|
||
print(f"🚀 开始处理涨停/炸板数据：({start_date} -> {end_date}]")
|
||
print(f"🧱 初始化并检查融合表 {COMBINED_TABLE} ...")
|
||
init_db(engine)
|
||
print("✅ 表结构、历史重复数据和唯一索引检查完成。")
|
||
|
||
requested_start_date = pd.to_datetime(start_date).date()
|
||
if requested_start_date < LOCAL_START_DATE:
|
||
print(f"ℹ️ 本地 stock_daily 暂按 {LOCAL_START_DATE} 起处理，已忽略更早日期。")
|
||
|
||
# 1. 获取基础数据和行情数据
|
||
query = text("""
|
||
SELECT
|
||
d.trade_date, d.symbol, d.ts_code, b.name,
|
||
d.high, d.close, d.pre_close, d.vol, b.market
|
||
FROM stock_daily d
|
||
JOIN stock_basic b
|
||
ON d.ts_code COLLATE utf8mb4_unicode_ci = b.ts_code COLLATE utf8mb4_unicode_ci
|
||
WHERE d.trade_date > :start_date
|
||
AND d.trade_date <= :end_date
|
||
AND d.trade_date >= :local_start_date
|
||
AND d.pre_close IS NOT NULL
|
||
ORDER BY d.trade_date ASC
|
||
""")
|
||
|
||
try:
|
||
print("📥 读取 stock_daily 与 stock_basic 数据 ...")
|
||
df = pd.read_sql(
|
||
query,
|
||
engine,
|
||
params={
|
||
"start_date": start_date,
|
||
"end_date": end_date,
|
||
"local_start_date": LOCAL_START_DATE,
|
||
},
|
||
)
|
||
if df.empty:
|
||
print(f"⚠️ 在指定区间 ({start_date} -> {end_date}] 内未获取到数据。")
|
||
return True
|
||
target_dates = df['trade_date'].astype(str).unique()
|
||
print(f"✅ 读取完成：{len(df)} 行行情，覆盖 {len(target_dates)} 个交易日。")
|
||
|
||
# 2. 涨停价计算规则
|
||
def calc_limit(row):
|
||
m, p = str(row['market']), float(row['pre_close'])
|
||
# 主板10%，创业板/科创板20%，北交所30%
|
||
ratio = 1.2 if ('创业板' in m or '科创板' in m) else (1.3 if '北交所' in m else 1.1)
|
||
return round(p * ratio, 2)
|
||
|
||
df['limit_price'] = df.apply(calc_limit, axis=1)
|
||
|
||
# 3. 筛选涨停(1)和炸板(2)
|
||
df['is_limit'] = (df['high'] == df['limit_price']) & (df['close'] == df['limit_price'])
|
||
df['is_broken'] = (df['high'] == df['limit_price']) & (df['close'] < df['limit_price'])
|
||
|
||
limit_df = df[df['is_limit']].copy()
|
||
limit_df['status'] = 1
|
||
|
||
broken_df = df[df['is_broken']].copy()
|
||
broken_df['status'] = 2
|
||
|
||
save_df = pd.concat([limit_df, broken_df])
|
||
print(
|
||
f"🧮 本地计算完成：涨停 {len(limit_df)} 条，炸板 {len(broken_df)} 条，"
|
||
f"合计 {len(save_df)} 条。"
|
||
)
|
||
|
||
tushare_df = fetch_tushare_limit_list(target_dates)
|
||
print("🔗 融合本地计算结果与 Tushare 数据 ...")
|
||
combined_df = build_combined_limit_data(df.copy(), save_df.copy(), tushare_df.copy())
|
||
validate_no_duplicate_frame_rows(combined_df, "融合待入库数据")
|
||
print(f"✅ 融合完成：待入库 {len(combined_df)} 条，重复键检查通过。")
|
||
|
||
with engine.begin() as conn:
|
||
delete_params = {
|
||
f"date_{idx}": value for idx, value in enumerate(target_dates)
|
||
}
|
||
placeholders = ", ".join(f":date_{idx}" for idx in range(len(target_dates)))
|
||
combined_delete_query = text(
|
||
f"DELETE FROM {COMBINED_TABLE} WHERE trade_date IN ({placeholders})"
|
||
)
|
||
print(f"🧹 清理 {len(target_dates)} 个交易日的旧融合数据 ...")
|
||
delete_result = conn.execute(combined_delete_query, delete_params)
|
||
print(f"✅ 旧数据清理完成：删除 {delete_result.rowcount} 条。")
|
||
|
||
if not combined_df.empty:
|
||
print(f"💾 写入 {COMBINED_TABLE}：{len(combined_df)} 条 ...")
|
||
combined_df.to_sql(
|
||
COMBINED_TABLE,
|
||
conn,
|
||
if_exists='append',
|
||
index=False,
|
||
chunksize=1000,
|
||
method="multi",
|
||
)
|
||
else:
|
||
print("ℹ️ 融合结果为空，本次不写入新数据。")
|
||
validate_no_duplicate_combined_rows(conn)
|
||
print("✅ 写入后数据库重复键校验通过。")
|
||
|
||
if not save_df.empty:
|
||
print(f"✅ 处理成功！涉及交易日: {len(target_dates)} 天")
|
||
print(f"📊 融合表 {COMBINED_TABLE} 入库: {len(combined_df)} 条记录。")
|
||
else:
|
||
print(f"ℹ️ 该时段内本地行情无涨停或炸板记录，已清理 {len(target_dates)} 个交易日的融合表旧结果。")
|
||
print(f"📊 融合表 {COMBINED_TABLE} 入库: {len(combined_df)} 条记录。")
|
||
return True
|
||
|
||
except Exception as e:
|
||
print(f"❌ 运行失败: {e}")
|
||
return False
|
||
|
||
|
||
def cleanup_combined_data():
|
||
engine = get_engine()
|
||
try:
|
||
init_db(engine, full_maintenance=True)
|
||
print(f"✅ {COMBINED_TABLE} 已完成历史重复数据清理、唯一索引检查和 calc_vol 回填。")
|
||
return True
|
||
except Exception as e:
|
||
print(f"❌ 清理失败: {e}")
|
||
return False
|
||
|
||
|
||
# ============================================================
|
||
# 用户配置区
|
||
# ============================================================
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser(description="计算并保存涨停/炸板数据")
|
||
parser.add_argument("--start-date", default="2026-05-11", help="开始日期，格式 YYYY-MM-DD，实际查询为大于该日期")
|
||
parser.add_argument("--end-date", default="2026-05-13", help="结束日期，格式 YYYY-MM-DD")
|
||
parser.add_argument(
|
||
"--cleanup-only",
|
||
action="store_true",
|
||
help="仅清理 stock_limit_combined_data 历史重复数据并从 stock_daily 回填 calc_vol",
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
if args.cleanup_only:
|
||
cleanup_combined_data()
|
||
raise SystemExit(0)
|
||
|
||
# 开始日期：计算从该日期的下一个交易日开始
|
||
START_DATE = args.start_date
|
||
|
||
# 结束日期
|
||
END_DATE = args.end_date
|
||
|
||
analyze_and_save_limits(START_DATE, END_DATE)
|