from __future__ import annotations import argparse import csv import json import subprocess import sys import time from concurrent.futures import ThreadPoolExecutor, as_completed from collections import Counter from datetime import datetime, timedelta from pathlib import Path from typing import Any import baostock as bs import pandas as pd PACKAGE_ROOT = Path(__file__).resolve().parents[1] CACHE_ROOT = PACKAGE_ROOT / "public_5m_proxy_baostock_cache" FETCH_LEDGER = PACKAGE_ROOT / "public_5m_proxy_baostock_fetch_inventory.csv" SOURCE_GAP = PACKAGE_ROOT / "three_day_high_minute_gap_inventory.csv" AUDIT_LEDGER = PACKAGE_ROOT / "three_day_high_5m_proxy_audit.csv" SUMMARY_PATH = PACKAGE_ROOT / "public_5m_proxy_summary.json" TEMP_ROOT = PACKAGE_ROOT / "public_5m_proxy_temp" LOCAL_DAILY_ROOT = Path("E:/quant/a_share_daily_front_20230101_20260508_complete/daily") FIELDS = "date,time,code,open,high,low,close,volume,amount,adjustflag" QUERY_WINDOW_DAYS = 7 def fnum(value: Any) -> float | None: text = str(value).strip() if not text: return None try: return float(text) except ValueError: return None def to_baostock_code(symbol: str) -> str | None: code, market = symbol.split(".") if market == "SH": return f"sh.{code}" if market == "SZ": return f"sz.{code}" return None def date_key(value: str) -> str: return value.replace("-", "")[:8] def iso_date(value: str) -> str: key = date_key(value) return f"{key[:4]}-{key[4:6]}-{key[6:8]}" def window_start_date(value: str) -> str: return (datetime.strptime(iso_date(value), "%Y-%m-%d") - timedelta(days=QUERY_WINDOW_DAYS)).strftime("%Y-%m-%d") def time_hhmm(value: str) -> str: text = str(value) if len(text) >= 12 and text[:8].isdigit(): return text[8:12] return "" def local_daily_high(symbol: str, trade_date: str, cache: dict[str, pd.DataFrame]) -> float | None: path = LOCAL_DAILY_ROOT / f"{symbol}.csv" if not path.exists(): return None if symbol not in cache: cache[symbol] = pd.read_csv(path, dtype=str, keep_default_na=False, encoding="utf-8-sig") df = cache[symbol] date_col = "trade_date" if "trade_date" in df.columns else "date" matched = df[df[date_col].map(date_key).eq(date_key(trade_date))] if len(matched) == 0 or "high" not in matched.columns: return None return fnum(matched.iloc[0]["high"]) def cache_path(symbol: str, trade_date: str) -> Path: return CACHE_ROOT / symbol.replace(".", "_") / f"{date_key(trade_date)}_5m_qfq_baostock.csv" def load_fetch_rows() -> list[dict[str, str]]: if FETCH_LEDGER.exists(): with FETCH_LEDGER.open("r", newline="", encoding="utf-8-sig") as f: return list(csv.DictReader(f)) gaps = pd.read_csv(SOURCE_GAP, dtype=str, keep_default_na=False, encoding="utf-8-sig") unique = gaps[["symbol", "observation_trade_date"]].drop_duplicates().sort_values(["symbol", "observation_trade_date"]) rows = [] for row in unique.to_dict("records"): rows.append( { "symbol": row["symbol"], "observation_trade_date": row["observation_trade_date"], "baostock_code": to_baostock_code(row["symbol"]) or "", "fetch_status": "PENDING", "row_count": "0", "cache_path": "", "error": "", } ) return rows def write_fetch_rows(rows: list[dict[str, str]]) -> None: keys = ["symbol", "observation_trade_date", "baostock_code", "fetch_status", "row_count", "cache_path", "error"] with FETCH_LEDGER.open("w", newline="", encoding="utf-8-sig") as f: writer = csv.DictWriter(f, fieldnames=keys, extrasaction="ignore") writer.writeheader() writer.writerows(rows) def fetch_one(symbol: str, trade_date: str) -> tuple[str, str, str, str]: bs_code = to_baostock_code(symbol) if bs_code is None: return "UNSUPPORTED_MARKET", "", "0", "Only SH/SZ supported by baostock." path = cache_path(symbol, trade_date) if path.exists() and path.stat().st_size > 0: try: row_count = len(pd.read_csv(path, dtype=str, keep_default_na=False, encoding="utf-8-sig")) except Exception: row_count = 0 if row_count > 0: return "FETCH_OK", str(path.relative_to(PACKAGE_ROOT)).replace("\\", "/"), str(row_count), "" login = bs.login() if login.error_code != "0": return "LOGIN_ERROR", "", "0", f"{login.error_code} {login.error_msg}" try: rs = bs.query_history_k_data_plus( bs_code, FIELDS, start_date=window_start_date(trade_date), end_date=iso_date(trade_date), frequency="5", adjustflag="2", ) rows: list[list[str]] = [] if rs.error_code == "0": while rs.next(): rows.append(rs.get_row_data()) df = pd.DataFrame(rows, columns=FIELDS.split(",")) if len(df) > 0 and "date" in df.columns: df = df[df["date"].map(date_key).eq(date_key(trade_date))] if len(df) > 0: path.parent.mkdir(parents=True, exist_ok=True) df.to_csv(path, index=False, encoding="utf-8-sig") return "FETCH_OK", str(path.relative_to(PACKAGE_ROOT)).replace("\\", "/"), str(len(df)), "" return "FETCH_EMPTY", "", "0", "" return "FETCH_ERROR", "", "0", f"{rs.error_code} {rs.error_msg}" finally: try: bs.logout() except Exception: pass def fetch_one_subprocess(symbol: str, trade_date: str, timeout_seconds: int) -> tuple[str, str, str, str]: TEMP_ROOT.mkdir(parents=True, exist_ok=True) output_path = TEMP_ROOT / f"{symbol.replace('.', '_')}_{date_key(trade_date)}_{int(time.time() * 1000)}.json" cmd = [ sys.executable, str(Path(__file__).resolve()), "--single-symbol", symbol, "--single-date", trade_date, "--single-output", str(output_path), ] try: result = subprocess.run( cmd, timeout=timeout_seconds, check=False, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True, ) except subprocess.TimeoutExpired: return "FETCH_TIMEOUT", "", "0", f"timeout after {timeout_seconds}s" if not output_path.exists(): stderr = (result.stderr or "").strip() return "FETCH_SUBPROCESS_ERROR", "", "0", stderr[:500] try: payload = json.loads(output_path.read_text(encoding="utf-8")) except json.JSONDecodeError as exc: return "FETCH_SUBPROCESS_ERROR", "", "0", f"invalid child output: {exc}" finally: try: output_path.unlink() except OSError: pass return ( str(payload.get("fetch_status", "")), str(payload.get("cache_path", "")), str(payload.get("row_count", "0")), str(payload.get("error", "")), ) def fetch_one_with_optional_timeout( symbol: str, trade_date: str, timeout_seconds: int ) -> tuple[str, str, str, str]: if timeout_seconds > 0: return fetch_one_subprocess(symbol, trade_date, timeout_seconds) return fetch_one(symbol, trade_date) def rebuild_audit_and_summary(fetch_rows: list[dict[str, str]]) -> dict[str, Any]: source = pd.read_csv(SOURCE_GAP, dtype=str, keep_default_na=False, encoding="utf-8-sig") fetch_by_key = {(row["symbol"], row["observation_trade_date"]): row for row in fetch_rows} audit_rows: list[dict[str, Any]] = [] daily_cache: dict[str, pd.DataFrame] = {} for row in source.to_dict("records"): symbol = row["symbol"] trade_date = row["observation_trade_date"] fetch = fetch_by_key.get((symbol, trade_date), {}) path_text = fetch.get("cache_path", "") raw_high_until_1040 = None raw_full_day_high = None local_current_day_high = None scale_to_local_daily = None high_until_1040 = None comparison_price_source = "" bar_count = 0 first_time = "" last_time = "" if fetch.get("fetch_status") == "FETCH_OK" and path_text: path = PACKAGE_ROOT / path_text if path.exists(): df = pd.read_csv(path, dtype=str, keep_default_na=False, encoding="utf-8-sig") if len(df) > 0: df["hhmm"] = df["time"].map(time_hhmm) df["high_num"] = pd.to_numeric(df["high"], errors="coerce") first_time = str(df["time"].iloc[0]) last_time = str(df["time"].iloc[-1]) before = df[df["hhmm"].le("1040")] bar_count = len(before) raw_full_day_high = float(df["high_num"].max()) if len(before) > 0: raw_high_until_1040 = float(before["high_num"].max()) local_current_day_high = local_daily_high(symbol, trade_date, daily_cache) if ( raw_high_until_1040 is not None and raw_full_day_high is not None and raw_full_day_high > 0 and local_current_day_high is not None ): scale_to_local_daily = local_current_day_high / raw_full_day_high high_until_1040 = raw_high_until_1040 * scale_to_local_daily comparison_price_source = "baostock_5m_scaled_to_local_daily_high" else: high_until_1040 = raw_high_until_1040 comparison_price_source = "baostock_5m_raw" h2 = fnum(row.get("d_minus_2_high_front")) h1 = fnum(row.get("d_minus_1_high_front")) not_rising = "" if h2 is not None and h1 is not None and high_until_1040 is not None: not_rising = not (h2 < h1 < high_until_1040) audit_rows.append( { "signal_id": row.get("signal_id", ""), "case_id": row.get("case_id", ""), "symbol": symbol, "observation_trade_date": trade_date, "candidate_time": row.get("candidate_time", ""), "minute_gap_status": row.get("minute_file_status", ""), "public_5m_source": "baostock query_history_k_data_plus frequency=5 adjustflag=2", "public_5m_fetch_status": fetch.get("fetch_status", ""), "public_5m_row_count": fetch.get("row_count", ""), "public_5m_cache_path": path_text, "public_5m_first_time": first_time, "public_5m_last_time": last_time, "public_5m_bar_count_until_1040": bar_count, "d_minus_2_high_front": row.get("d_minus_2_high_front", ""), "d_minus_1_high_front": row.get("d_minus_1_high_front", ""), "current_high_until_1040_5m_qfq_raw": raw_high_until_1040 if raw_high_until_1040 is not None else "", "public_5m_full_day_high_qfq_raw": raw_full_day_high if raw_full_day_high is not None else "", "local_daily_observation_high_front": local_current_day_high if local_current_day_high is not None else "", "public_5m_to_local_daily_scale": scale_to_local_daily if scale_to_local_daily is not None else "", "current_high_until_1040_5m_daily_scaled_front": high_until_1040 if high_until_1040 is not None else "", "comparison_price_source": comparison_price_source, "three_highs_not_strictly_rising_5m_proxy": not_rising, "verification_status_5m_proxy": "PASS_5M_PROXY" if not_rising is True else ("FAIL_5M_PROXY" if not_rising is False else "NO_5M_DATA"), "code_evidence_reason_cn": row.get("code_evidence_reason_cn", ""), } ) keys: list[str] = [] seen: set[str] = set() for row in audit_rows: for key in row: if key not in seen: seen.add(key) keys.append(key) with AUDIT_LEDGER.open("w", newline="", encoding="utf-8-sig") as f: writer = csv.DictWriter(f, fieldnames=keys, extrasaction="ignore") writer.writeheader() writer.writerows(audit_rows) summary = { "generated_at": datetime.now().isoformat(timespec="seconds"), "input_rows": len(source), "unique_symbol_dates_requested": len(fetch_rows), "fetch_status_counts": dict(Counter(row.get("fetch_status", "") for row in fetch_rows)), "audit_status_counts": dict(Counter(str(row["verification_status_5m_proxy"]) for row in audit_rows)), "boundaries": [ "This is a public 5-minute K-line proxy, not a replacement for original 1-minute evidence.", "Bars are treated as period-ending timestamps; bars with hhmm <= 10:40 are included for the 10:40 high proxy.", "Baostock adjustflag=2 is used as a front-adjusted proxy; raw Baostock prices are scaled to the local daily front-adjusted high before comparing with local daily highs when local daily data is available.", "The proxy is used only to validate the direction of SELL_THREE_DAY_HIGH_NOT_RISING.", ], } SUMMARY_PATH.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8") return summary def write_manifest() -> None: rows = [] for path in sorted(PACKAGE_ROOT.rglob("*")): if path.is_file(): rows.append({"path": str(path.relative_to(PACKAGE_ROOT)).replace("\\", "/"), "bytes": path.stat().st_size}) with (PACKAGE_ROOT / "manifest.csv").open("w", newline="", encoding="utf-8-sig") as f: writer = csv.DictWriter(f, fieldnames=["path", "bytes"]) writer.writeheader() writer.writerows(rows) def write_progress(path_text: str, payload: dict[str, Any]) -> None: if not path_text: return path = Path(path_text) if not path.is_absolute(): path = PACKAGE_ROOT / path path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8") def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--limit", type=int, default=20) parser.add_argument("--sleep", type=float, default=0.5) parser.add_argument("--progress-path", default="") parser.add_argument("--done-path", default="") parser.add_argument("--per-request-timeout", type=int, default=0) parser.add_argument("--workers", type=int, default=1) parser.add_argument("--retry-status", default="") parser.add_argument("--single-symbol", default="") parser.add_argument("--single-date", default="") parser.add_argument("--single-output", default="") args = parser.parse_args() if args.single_symbol: status, cache, count, error = fetch_one(args.single_symbol, args.single_date) payload = { "symbol": args.single_symbol, "observation_trade_date": args.single_date, "fetch_status": status, "cache_path": cache, "row_count": count, "error": error, } if args.single_output: output_path = Path(args.single_output) output_path.parent.mkdir(parents=True, exist_ok=True) output_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8") else: print(json.dumps(payload, ensure_ascii=False, indent=2)) return rows = load_fetch_rows() attempted = 0 retry_statuses = {item.strip() for item in args.retry_status.split(",") if item.strip()} if retry_statuses: pending_rows = [row for row in rows if row.get("fetch_status") in retry_statuses][: args.limit] else: pending_rows = [row for row in rows if row.get("fetch_status") != "FETCH_OK"][: args.limit] def update_row(row: dict[str, str], result: tuple[str, str, str, str]) -> None: nonlocal attempted status, cache, count, error = result row["baostock_code"] = to_baostock_code(row["symbol"]) or "" row["fetch_status"] = status row["row_count"] = count row["cache_path"] = cache row["error"] = error attempted += 1 write_fetch_rows(rows) write_progress( args.progress_path, { "generated_at": datetime.now().isoformat(timespec="seconds"), "status": "RUNNING", "attempted": attempted, "limit": args.limit, "workers": max(args.workers, 1), "current_symbol": row["symbol"], "current_trade_date": row["observation_trade_date"], "current_fetch_status": status, "fetch_status_counts": dict(Counter(item.get("fetch_status", "") for item in rows)), }, ) if args.sleep > 0: time.sleep(args.sleep) if max(args.workers, 1) > 1 and pending_rows: with ThreadPoolExecutor(max_workers=max(args.workers, 1)) as executor: futures = { executor.submit( fetch_one_with_optional_timeout, row["symbol"], row["observation_trade_date"], args.per_request_timeout, ): row for row in pending_rows } for future in as_completed(futures): row = futures[future] try: result = future.result() except Exception as exc: result = ("FETCH_WORKER_ERROR", "", "0", str(exc)) update_row(row, result) else: for row in pending_rows: result = fetch_one_with_optional_timeout( row["symbol"], row["observation_trade_date"], args.per_request_timeout, ) update_row(row, result) summary = rebuild_audit_and_summary(rows) write_manifest() final_payload = {"attempted": attempted, **summary} write_progress(args.progress_path, {"status": "FINISHED", **final_payload}) write_progress(args.done_path, {"status": "FINISHED", **final_payload}) print(json.dumps(final_payload, ensure_ascii=False, indent=2)) if __name__ == "__main__": main()