#!/usr/bin/env python3 """生成 / 校验缠论流水线 golden（行为冻结基线）。""" from __future__ import annotations import argparse import json import sys from pathlib import Path import numpy as np import pandas as pd ROOT = Path(__file__).resolve().parents[1] if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) def make_ohlcv(n: int = 400, seed: int = 42) -> pd.DataFrame: rng = np.random.default_rng(seed) dates = pd.date_range("2024-01-01", periods=n, freq="5min", tz="UTC") # 随机游走 + 波动，保证高低开收合法 rets = rng.normal(0, 0.002, size=n) close = 100 * np.exp(np.cumsum(rets)) open_ = np.roll(close, 1) open_[0] = close[0] spread = np.abs(rng.normal(0, 0.0015, size=n)) * close high = np.maximum(open_, close) + spread low = np.minimum(open_, close) - spread volume = rng.uniform(100, 1000, size=n) return pd.DataFrame( { "date": dates, "open": open_, "high": high, "low": low, "close": close, "volume": volume, } ) def _enum_name(v): if v is None: return None name = getattr(v, "name", None) if name: return name s = str(v) return s.split(".")[-1] if "." in s else s def _t(obj, attr="time"): t = getattr(obj, attr, None) if t is None: return None if hasattr(t, "strftime"): return t.strftime("%Y-%m-%d %H:%M:%S") return str(t) def _f(v, nd=10): if v is None: return None try: return round(float(v), nd) except (TypeError, ValueError): return None def serialize_pipeline(tf) -> dict: bi_list = [ { "idx": getattr(b, "idx", i), "dir": _enum_name(getattr(b, "dir", None)), "high": _f(b.get_high() if hasattr(b, "get_high") else getattr(b, "high", 0)), "low": _f(b.get_low() if hasattr(b, "get_low") else getattr(b, "low", 0)), "is_sure": bool(getattr(b, "is_sure", True)), "begin": _t(getattr(b, "begin_klc", None), "end_time") or _t(getattr(b, "begin_klc", None), "time"), "end": _t(getattr(b, "end_klc", None), "end_time") or _t(getattr(b, "end_klc", None), "time"), } for i, b in enumerate(getattr(tf, "bi_list", []) or []) ] seg_list = [ { "idx": getattr(s, "idx", i), "dir": _enum_name(getattr(s, "dir", None)), "is_sure": bool(getattr(s, "is_sure", True)), "high": _f(s.get_high()) if hasattr(s, "get_high") else None, "low": _f(s.get_low()) if hasattr(s, "get_low") else None, } for i, s in enumerate(getattr(tf, "seg_list", []) or []) ] zs_list = [ { "idx": getattr(z, "idx", i), "dir": _enum_name(getattr(z, "dir", None)), "zg": _f(getattr(z, "zg", 0) or 0), "zd": _f(getattr(z, "zd", 0) or 0), "is_sure": bool(getattr(z, "is_sure", True)), } for i, z in enumerate(getattr(tf, "zs_list", []) or []) ] bsp_list = [] for i, p in enumerate(getattr(tf, "bsp_list", []) or []): bsp_list.append( { "idx": getattr(p, "idx", i), "type": _enum_name(getattr(p, "type", None) or getattr(p, "bsp_type", None)), "dir": _enum_name(getattr(p, "dir", None)), "price": _f(getattr(p, "price", 0) or getattr(p, "val", 0) or 0), } ) return { "counts": { "klu": len(getattr(tf, "klu_list", []) or []), "klc": len(getattr(tf, "klc_list", []) or []), "bi": len(bi_list), "seg": len(seg_list), "zs": len(zs_list), "bsp": len(bsp_list), }, "bi_list": bi_list, "seg_list": seg_list, "zs_list": zs_list, "bsp_list": bsp_list, } def analyze_contract_keys() -> list: """文档化 /api/analyze 主周期关键字段（契约冒烟用）。""" return sorted( [ "timezone", "kline_data", "klc_list", "bi_list", "uncompleted_bi_list", "seg_list", "uncompleted_seg_list", "zs_list", "uncompleted_zs_list", "bi_zs_list", "bsp_list", "klc_fx_info", "macd", "chan_macd", "klc_trend", ] ) def run_pipeline(df: pd.DataFrame): """与 web.analyze_chan 同序，避免依赖 TF_DF.__init__ 历史缺口。""" from chanlun import TF_DF chan = TF_DF() df = chan.add_indicators(df.copy()) klu_list = chan.get_kl_data(df) klc_list = chan.get_klc_list(klu_list) bi_list = chan.cal_bi_list(klc_list) seg_list = chan.get_seg_list(bi_list) zs_list = chan.calculate_seg_zs(seg_list) bi_zs_list = chan.cal_bi_zs(seg_list) bsp_list = chan.find_all_bsp(bi_list, bi_zs_list) if bi_zs_list else [] # 挂到伪对象供 serialize_pipeline 使用 chan.klu_list = klu_list chan.klc_list = klc_list chan.bi_list = bi_list chan.seg_list = seg_list chan.zs_list = zs_list chan.bsp_list = bsp_list chan.bi_zs_list = bi_zs_list return chan def main(): parser = argparse.ArgumentParser() parser.add_argument("--generate", action="store_true") parser.add_argument("--check", action="store_true") parser.add_argument( "--fixture-dir", type=Path, default=ROOT / "tests" / "fixtures", ) args = parser.parse_args() args.fixture_dir.mkdir(parents=True, exist_ok=True) ohlcv_path = args.fixture_dir / "ohlcv_5m.csv" golden_path = args.fixture_dir / "golden_pipeline.json" keys_path = args.fixture_dir / "analyze_contract_keys.json" if args.generate: df = make_ohlcv() df.to_csv(ohlcv_path, index=False) # 与 --check 同一路径：经 CSV 往返，避免浮点路径差异 df = pd.read_csv(ohlcv_path, parse_dates=["date"]) tf = run_pipeline(df) payload = serialize_pipeline(tf) golden_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8") keys_path.write_text(json.dumps(analyze_contract_keys(), indent=2), encoding="utf-8") print(f"wrote {ohlcv_path}") print(f"wrote {golden_path} counts={payload['counts']}") return 0 if args.check: df = pd.read_csv(ohlcv_path, parse_dates=["date"]) tf = run_pipeline(df) actual = serialize_pipeline(tf) expected = json.loads(golden_path.read_text(encoding="utf-8")) if actual != expected: print("GOLDEN MISMATCH") print("expected counts", expected.get("counts")) print("actual counts", actual.get("counts")) diff_path = args.fixture_dir / "golden_pipeline.actual.json" diff_path.write_text(json.dumps(actual, ensure_ascii=False, indent=2), encoding="utf-8") print(f"wrote actual to {diff_path}") return 1 print("GOLDEN OK", actual["counts"]) return 0 parser.print_help() return 2 if __name__ == "__main__": raise SystemExit(main())