108 lines
4.5 KiB
Python
108 lines
4.5 KiB
Python
"""共用：股票码规范化、覆盖池、幂等写因子表、注册 factor_metadata（自适应列）。"""
|
||
import json
|
||
|
||
import pandas as pd
|
||
|
||
import db
|
||
|
||
|
||
def to_prefix(ts_code: str) -> str:
|
||
"""600000.SH -> SH600000；已是前缀式(SZ002625)或纯代码则原样返回。"""
|
||
s = str(ts_code).strip()
|
||
if "." in s:
|
||
num, exch = s.split(".", 1)
|
||
return f"{exch.upper()}{num}"
|
||
return s
|
||
|
||
|
||
def load_universe() -> set[str]:
|
||
"""KG 覆盖池 ts_code 集合（600000.SH 形态），来自基座视图 v_factor_universe。"""
|
||
df = db.read_pg("SELECT ts_code FROM v_factor_universe")
|
||
return set(df["ts_code"].astype(str).str.strip())
|
||
|
||
|
||
def trading_days(start: str, end: str) -> list:
|
||
"""目标区间交易日历——用热度表(153，日频、覆盖广)的 distinct trade_date 近似。"""
|
||
df = db.read_mysql("heat",
|
||
"SELECT DISTINCT trade_date FROM stock_fund_heat_scores "
|
||
"WHERE trade_date BETWEEN %s AND %s ORDER BY trade_date",
|
||
(start, end))
|
||
return list(pd.to_datetime(df["trade_date"]))
|
||
|
||
|
||
_CREATE_FACTOR_TABLE = """
|
||
CREATE TABLE IF NOT EXISTS {t} (
|
||
trade_date DATE NOT NULL,
|
||
stock_code VARCHAR(15) NOT NULL,
|
||
factor_value DOUBLE,
|
||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||
PRIMARY KEY (trade_date, stock_code),
|
||
KEY idx_stock (stock_code)
|
||
) ENGINE=InnoDB
|
||
"""
|
||
|
||
|
||
def ensure_table(table: str) -> None:
|
||
with db.factor_conn() as conn:
|
||
with conn.cursor() as cur:
|
||
cur.execute(_CREATE_FACTOR_TABLE.format(t=table))
|
||
conn.commit()
|
||
|
||
|
||
def write_factor(table: str, df: pd.DataFrame, mode: str) -> None:
|
||
"""df 列 [trade_date, stock_code, factor_value] → 幂等写因子表。
|
||
daily/history 都是「删涉及日期区间 → 批插」。stock_code 统一转前缀式。"""
|
||
if df is None or df.empty:
|
||
print(f" {table}: 无数据（跳过）")
|
||
return
|
||
df = df.dropna(subset=["trade_date", "stock_code", "factor_value"]).copy()
|
||
df["stock_code"] = df["stock_code"].map(to_prefix)
|
||
df["trade_date"] = pd.to_datetime(df["trade_date"]).dt.date
|
||
df = df.drop_duplicates(["trade_date", "stock_code"], keep="last")
|
||
if df.empty:
|
||
print(f" {table}: 清洗后无数据")
|
||
return
|
||
dmin, dmax = df["trade_date"].min(), df["trade_date"].max()
|
||
ensure_table(table)
|
||
rows = list(df[["trade_date", "stock_code", "factor_value"]]
|
||
.itertuples(index=False, name=None))
|
||
with db.factor_conn() as conn:
|
||
with conn.cursor() as cur:
|
||
cur.execute(f"DELETE FROM {table} WHERE trade_date BETWEEN %s AND %s", (dmin, dmax))
|
||
cur.executemany(
|
||
f"INSERT INTO {table} (trade_date, stock_code, factor_value) VALUES (%s,%s,%s)",
|
||
rows)
|
||
conn.commit()
|
||
print(f" {table}: 写入 {len(rows)} 行, 日期 {dmin}~{dmax}")
|
||
|
||
|
||
def register(factor_code, display_name, table, category, desc,
|
||
author="akg-factor-bridge") -> None:
|
||
"""注册 factor_metadata。自适应实际存在的列（避免猜死 schema）——
|
||
列名以线上 factor_metadata 为准，本函数只写它有的列。"""
|
||
values = {
|
||
"factor_code": factor_code, "display_name": display_name,
|
||
"target_ds_name": "ds_a", "target_table_name": table,
|
||
"author": author, "description": desc,
|
||
"category": json.dumps(category, ensure_ascii=False),
|
||
"frequency": "daily", "status": "active", "factor_type": "single",
|
||
}
|
||
with db.factor_conn() as conn:
|
||
with conn.cursor() as cur:
|
||
cur.execute("SELECT column_name FROM information_schema.columns "
|
||
"WHERE table_schema=DATABASE() AND table_name='factor_metadata'")
|
||
cols = {r[0] for r in cur.fetchall()}
|
||
use = [k for k in values if k in cols]
|
||
if "factor_code" not in use:
|
||
raise RuntimeError("factor_metadata 无 factor_code 列？请核实线上 schema")
|
||
ph = ",".join(["%s"] * len(use))
|
||
upd = ",".join(f"{k}=VALUES({k})" for k in use if k != "factor_code")
|
||
cur.execute(
|
||
f"INSERT INTO factor_metadata ({','.join(use)}) VALUES ({ph}) "
|
||
f"ON DUPLICATE KEY UPDATE {upd}", [values[k] for k in use])
|
||
conn.commit()
|
||
print(f" 注册 {factor_code} -> {table}（写入列: {sorted(use)}）")
|
||
if "factor_type" not in cols:
|
||
print(" ⚠️ factor_metadata 无 factor_type 列——平台 /mining/factors/all 按 "
|
||
"factor_type IN('single','multiple') 过滤，缺列可能查不到本因子，请核实。")
|