161 lines
7.3 KiB
Python
161 lines
7.3 KiB
Python
"""通达信「沪深京日线数据完整包」全量导入 candles 表。
|
||
|
||
用法（在 backend 目录下）：
|
||
uv run python scripts/import_tdx_day.py C:/Users/cirry/Downloads/hsjday [symbol ...]
|
||
# symbol 为可选的 6 位代码过滤（如 000001 002671），只重导这些标的
|
||
uv run python scripts/import_tdx_day.py <目录> --no-clear
|
||
# --no-clear：不清空任何行，纯 upsert（用于给已导入的底座回补 amount 成交额）
|
||
|
||
- 解析 vipdoc 的 .day 二进制文件（每条 32 字节）：
|
||
日期(YYYYMMDD) 开 高 低 收（×100） 成交额(元, float32) 成交量(股) 保留
|
||
- 只导入 stock_basic 里登记的股票（自动排除指数/基金/可转债/回购）；
|
||
sh000001(上证指数) 与 sz000001(平安银行) 这类代码冲突也由此化解。
|
||
- 价格为**不复权**：全量模式导入前清空已有的非 DEMO 行情；指定 symbol 过滤时
|
||
只清空这些标的（用于修复被复权口径污染的个别股票），其余不动。
|
||
- amount 为 TDX 原生 float32（元），精度 ~6 位有效数字，展示用途足够；
|
||
ON CONFLICT 时仅更新 amount 列，不动 OHLCV/turnover（避免与换手率回补互相干扰）。
|
||
- 写入用 asyncpg execute_many + ON CONFLICT DO UPDATE，可重复执行（幂等）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import asyncio
|
||
import struct
|
||
import sys
|
||
import time
|
||
from pathlib import Path
|
||
|
||
import asyncpg
|
||
|
||
# .env 里的 DATABASE_URL 是 SQLAlchemy 格式，asyncpg 需要 libpq 格式
|
||
DEFAULT_URL = "postgresql://postgres:postgres@localhost:5432/stock"
|
||
BATCH = 20_000 # 每批 upsert 行数
|
||
|
||
|
||
def load_db_url() -> str:
|
||
env = Path(__file__).resolve().parent.parent / ".env"
|
||
if env.exists():
|
||
for line in env.read_text(encoding="utf-8").splitlines():
|
||
line = line.strip()
|
||
if line.startswith("DATABASE_URL=postgresql+asyncpg://"):
|
||
return "postgresql://" + line.split("://", 1)[1]
|
||
return DEFAULT_URL
|
||
|
||
|
||
def parse_day_file(path: Path) -> list[tuple[int, float, float, float, float, float, float]]:
|
||
"""解析单个 .day 文件 -> [(date, open, high, low, close, volume(股), amount(元)), ...]"""
|
||
raw = path.read_bytes()
|
||
unpack = struct.Struct("<IIIIIfII").unpack_from
|
||
out = []
|
||
for i in range(len(raw) // 32):
|
||
date, o, h, l, c, amount, vol, _reserved = unpack(raw, i * 32)
|
||
out.append((date, o / 100.0, h / 100.0, l / 100.0, c / 100.0, float(vol), float(amount)))
|
||
return out
|
||
|
||
|
||
async def main(root: Path, symbols: list[str] | None = None, no_clear: bool = False) -> None:
|
||
if not root.exists():
|
||
sys.exit(f"目录不存在: {root}")
|
||
|
||
conn = await asyncpg.connect(load_db_url())
|
||
try:
|
||
# 股票清单：ts_code 形如 000001.SZ，用于过滤指数/基金/转债
|
||
rows = await conn.fetch("SELECT ts_code, symbol FROM stock_basic WHERE list_status = 'L'")
|
||
by_exchange: dict[str, set[str]] = {"sh": set(), "sz": set(), "bj": set()}
|
||
for r in rows:
|
||
suffix = r["ts_code"].split(".")[-1].lower() # SH/SZ/BJ -> sh/sz/bj
|
||
if suffix in by_exchange:
|
||
by_exchange[suffix].add(r["symbol"])
|
||
print(f"stock_basic 在市股票: " + ", ".join(f"{k}={len(v)}" for k, v in by_exchange.items()))
|
||
|
||
files = sorted(root.glob("*/lday/*.day"))
|
||
print(f"发现 .day 文件: {len(files)} 个")
|
||
|
||
if no_clear:
|
||
print("--no-clear：不清空任何行，纯 upsert 回补 amount")
|
||
elif symbols:
|
||
# 清空旧行情（保留 DEMO 合成数据），避免 qfq/不复权混用；
|
||
# 带 symbol 过滤时只清空目标标的（修复个别被污染的股票，不动其余底座）
|
||
deleted = await conn.execute(
|
||
"DELETE FROM candles WHERE symbol = ANY($1)", symbols
|
||
)
|
||
print(f"清空目标标的 {symbols}: {deleted}")
|
||
keep = set(symbols)
|
||
files = [p for p in files if p.name[2:8] in keep]
|
||
print(f"过滤后待导入 .day 文件: {len(files)} 个")
|
||
else:
|
||
deleted = await conn.execute("DELETE FROM candles WHERE symbol <> 'DEMO'")
|
||
print(f"清空旧行情: {deleted}")
|
||
|
||
if no_clear:
|
||
# 回填模式：只写 amount，不动 OHLCV/turnover（底座已就位，避免全表重写）
|
||
upsert_sql = """
|
||
INSERT INTO candles (symbol, timeframe, ts, open, high, low, close, volume, amount)
|
||
VALUES ($1, '1d', to_timestamp($2::text, 'YYYYMMDD')::timestamp, $3, $4, $5, $6, $7, $8)
|
||
ON CONFLICT (symbol, timeframe, ts) DO UPDATE
|
||
SET amount = EXCLUDED.amount
|
||
"""
|
||
else:
|
||
upsert_sql = """
|
||
INSERT INTO candles (symbol, timeframe, ts, open, high, low, close, volume, amount)
|
||
VALUES ($1, '1d', to_timestamp($2::text, 'YYYYMMDD')::timestamp, $3, $4, $5, $6, $7, $8)
|
||
ON CONFLICT (symbol, timeframe, ts) DO UPDATE
|
||
SET open = EXCLUDED.open, high = EXCLUDED.high, low = EXCLUDED.low,
|
||
close = EXCLUDED.close, volume = EXCLUDED.volume, amount = EXCLUDED.amount
|
||
"""
|
||
|
||
t0 = time.time()
|
||
total_stocks = 0
|
||
skipped = 0
|
||
batch: list[tuple] = []
|
||
|
||
rows_done = 0
|
||
|
||
async def flush() -> None:
|
||
nonlocal batch, rows_done
|
||
if batch:
|
||
await conn.executemany(upsert_sql, batch)
|
||
rows_done += len(batch)
|
||
batch = []
|
||
|
||
for n, path in enumerate(files, 1):
|
||
market = path.name[:2].lower() # sh / sz / bj
|
||
code = path.name[2:8]
|
||
if code not in by_exchange.get(market, set()):
|
||
skipped += 1
|
||
continue
|
||
for date, o, h, l, c, v, amount in parse_day_file(path):
|
||
batch.append((code, str(date), o, h, l, c, v, amount))
|
||
total_stocks += 1
|
||
if len(batch) >= BATCH:
|
||
await flush()
|
||
if n % 500 == 0:
|
||
elapsed = time.time() - t0
|
||
print(f" 进度 {n}/{len(files)} 文件, 已入库 {total_stocks} 只股票, "
|
||
f"{rows_done + len(batch):,} 行, {elapsed:.0f}s")
|
||
await flush()
|
||
|
||
cnt = await conn.fetchval("SELECT count(*) FROM candles WHERE symbol <> 'DEMO'")
|
||
span = await conn.fetchrow(
|
||
"SELECT min(ts) AS lo, max(ts) AS hi FROM candles WHERE symbol <> 'DEMO'"
|
||
)
|
||
with_amt = await conn.fetchval(
|
||
"SELECT count(*) FROM candles WHERE symbol <> 'DEMO' AND amount IS NOT NULL"
|
||
)
|
||
print(f"\n完成: {total_stocks} 只股票, {cnt:,} 行日线, "
|
||
f"范围 {span['lo']:%Y-%m-%d} ~ {span['hi']:%Y-%m-%d}, "
|
||
f"含成交额 {with_amt:,} 行, "
|
||
f"跳过非股票文件 {skipped} 个, 耗时 {time.time() - t0:.0f}s")
|
||
finally:
|
||
await conn.close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
ap = argparse.ArgumentParser(description="TDX 沪深京日线全量导入 candles")
|
||
ap.add_argument("root", help="hsjday 目录（其下 */lday/*.day）")
|
||
ap.add_argument("symbols", nargs="*", help="可选的 6 位代码过滤")
|
||
ap.add_argument("--no-clear", action="store_true",
|
||
help="不清空任何行，纯 upsert（amount 回补模式）")
|
||
a = ap.parse_args()
|
||
asyncio.run(main(Path(a.root), a.symbols or None, a.no_clear))
|