hamuna-quant-cli 0.1.0.dev93__py3-none-any.whl

hamuna_quant_cli/prebuilt_downloader.py

@@ -0,0 +1,263 @@
+"""ADR-0040 Phase C — prebuilt 数据集下载器 (v2 自有版本, B6).
+
+v2 自足版 (2026-08-19 修复): 原 wholesale copy 引用了包根不存在的
+`cache` / `server_client` / `market` 模块 (`from . import cache` 挂 =
+dataset list/fetch 全废)。本版改为直用 v2 runtime:
+  - runtime.cache          — bundle/single 落盘路径 (~/.hamuna/data_cache/datasets/)
+  - runtime.server_client  — list_datasets / s3_credentials (加密 S3 凭据)
+  - runtime.s3client       — HKDF 解密 + boto3 全头签名下载 (COS 反代 presign 403)
+
+三层下载:
+  list_available()      — 调 server `list_datasets` 枚举内置预构建
+  download_bundle(uni)   — 加密凭据 + boto3 直连 → 落 __bundle__<uni>_<dt>.parquet
+  download_single(sym)   — 从所属 bundle 切片 → 落 <bare>_<dt>_fq1.parquet
+                          (bundle 缺失则先整包下载; server 无 per-symbol artifact)
+
+CLI:
+  hamuna_quant_cli dataset list
+  hamuna_quant_cli dataset fetch --symbols 600000,600036 [--start] [--end]
+  (bundle/all 入口保留: python -m hamuna_quant_cli.prebuilt_downloader fetch bundle all_a)
+
+隔离:
+  - 不重新实现 S3 / 容维 client; 全部走 v2 包内 runtime
+  - 失败统一抛 RuntimeError, CLI 出口码 = 2 (与 hamuna_quant_cli 一致: 数据缺失)
+  - 新依赖: boto3 + cryptography (pyproject 已声明, S3 签名下载必需)
+
+Ponytail ceiling: 只支持日线 (period='1d' 唯一, 与 KNOWN_PREBUILT 对齐).
+升级: 加 5m/15m 需先扩 _PREBUILT_TYPE_CODE 注册 + server 端预构建 + 容维 minute 权限.
+"""
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+from .runtime import cache as _cache
+from .runtime.http_client import _load_token
+from .runtime.server_client import list_datasets, s3_credentials
+from .runtime.s3client import decrypt_s3_params, download_to
+
+# dt → server data_type_code (server 目前仅日线预构建; 5m 有数据集但 resolver 锁 1d)
+_PREBUILT_TYPE_CODE: dict[str, str] = {'D': 'kline_d'}
+
+# universe → 覆盖前缀 (与 prebuilt_resolver.KNOWN_PREBUILT 的 bundle_prefix_predicate 同步)
+_BUNDLE_PREFIXES: dict[str, tuple[str, ...]] = {
+    'all_a': ('60', '68', '00', '30', '40', '43', '83', '87', '88'),
+    'etf': ('51', '56', '15', '16'),
+    'convertible_bond': ('11', '12'),
+}
+
+# single parquet 写入列 (bundle 切片; resolver._read_single_with_pushdown 认 time/OHLCV)
+_SINGLE_COLUMNS = ('time', 'stockCode', 'open', 'high', 'low', 'close', 'volume', 'amount')
+
+
+def _normalize_symbol(symbol: str) -> str:
+    """裸码 → 带后缀 (与 akquant_schema_adapter.normalize_symbol 同规则, 避免循环 import)."""
+    s = symbol.strip().upper()
+    if s.endswith(('.SH', '.SZ', '.BJ')):
+        return symbol
+    if len(s) > 2 and s[:2] in ('SH', 'SZ', 'BJ'):
+        s = s[2:]
+    if len(s) != 6 or not s.isdigit():
+        return symbol
+    if s.startswith(('60', '68', '11', '13')):
+        return f'{s}.SH'
+    return f'{s}.SZ'
+
+
+def _bundle_for(symbol: str) -> str | None:
+    """带后缀 symbol → 所属 bundle universe 名 (前缀匹配; 无 → None)."""
+    bare = symbol.split('.')[0]
+    for uni, prefixes in _BUNDLE_PREFIXES.items():
+        if any(bare.startswith(p) for p in prefixes):
+            return uni
+    return None
+
+
+# ---- API ----
+
+def list_available() -> list[dict]:
+    """GET /api/v1/datasets → server 内置预构建列表.
+
+    每条 {universe_code, data_type_code, builtin, enabled, ...}.
+    返 [] 表示请求失败 / 无内置 — CLI 提示"无预构建".
+    """
+    try:
+        return list_datasets()
+    except Exception as e:  # noqa: BLE001 — 与 v1 同语义: 请求失败 → [] 引导
+        print(f'[warn] list_datasets 失败: {type(e).__name__}: {str(e)[:160]}', file=sys.stderr)
+        return []
+
+
+def download_bundle(universe: str, dt: str = 'D') -> Path:
+    """下载 server 预构建整包 → 落 __bundle__<uni>_<dt>.parquet.
+
+    走 s3_credentials (加密凭据) + s3client (boto3 全头签名直连) — COS 反代
+    presign 403, 这是唯一可靠路径。本地已存在 → 直接复用 (无 TTL, 历史不变).
+
+    Raises:
+      RuntimeError: 整包下载失败 / 非预构建类型 / server 无此内置
+    """
+    type_code = _PREBUILT_TYPE_CODE.get(dt)
+    if type_code is None:
+        raise RuntimeError(
+            f'dt={dt!r} 无预构建 (仅日线 D). 已知: {list(_PREBUILT_TYPE_CODE.keys())}'
+        )
+    path = _cache.bundle_path(universe, dt)
+    if path.exists() and path.stat().st_size > 0:
+        return path  # 已落盘复用 (无 TTL, 历史不变)
+
+    available = [d for d in list_available()
+                 if d.get('universe_code') == universe
+                 and d.get('data_type_code') == type_code
+                 and d.get('builtin') and d.get('enabled')]
+    if not available:
+        raise RuntimeError(
+            f'server 无内置预构建 {universe}/{dt} (或 builtin/enabled 未开)'
+        )
+    did = available[0].get('id') or available[0].get('_id')
+    if not did:
+        raise RuntimeError(f'dataset {universe}/{dt} 缺 id: {available[0]}')
+
+    # S3 key 布局: <datasetID>/<universe>/<type_code>/_merged/all.parquet
+    key = f'{did}/{universe}/{type_code}/_merged/all.parquet'
+    try:
+        env = s3_credentials(did, key)
+        params = decrypt_s3_params(_load_token(), did, env)
+        download_to(params, path)
+    except Exception as e:  # noqa: BLE001 — 失败原因必须可见, 不静默回退
+        raise RuntimeError(f'整包下载失败 {universe}/{type_code}: {type(e).__name__}: {e}') from e
+    return path
+
+
+def download_single(symbol: str, period: str = '1d',
+                     start: str | None = None,
+                     end: str | None = None) -> Path:
+    """下载单股数据集 → 落 <bare>_<dt>_fq1.parquet.
+
+    从所属 bundle 切片 (server 无 per-symbol artifact; bundle 覆盖 2019~今,
+    全量历史, start/end 不截断 — 语义 = "确保该股 dataset 可用").
+    bundle 缺失 → 先整包下载 (一次 422MB, 之后全部 single 复用).
+
+    Raises:
+      RuntimeError: bundle 不含该股 / 切片为空 (次新股/停牌/数据不可达)
+    """
+    if period != '1d':
+        raise RuntimeError(
+            f'period={period!r} 仅支持 1d (与 KNOWN_PREBUILT 对齐); '
+            f'非日线 dataset build 走 strategy_cli dataset --config'
+        )
+    full = _normalize_symbol(symbol)
+    bare = full.split('.')[0]
+    universe = _bundle_for(full)
+    if universe is None:
+        raise RuntimeError(f'{full} 不属于任何预构建 universe (all_a/etf/convertible_bond)')
+
+    bundle = download_bundle(universe, 'D')
+    import pyarrow.parquet as pq
+    try:
+        table = pq.read_table(
+            bundle, columns=list(_SINGLE_COLUMNS),
+            filters=[('stockCode', '=', full)],
+        )
+    except Exception as e:  # noqa: BLE001
+        raise RuntimeError(f'bundle 读取失败 ({bundle}): {e}') from e
+    if table.num_rows == 0:
+        raise RuntimeError(
+            f'bundle {universe} 不含 {full} (可能退市/停更/不在成分池)'
+        )
+    out = _cache._dataset_path(full, 'D', 1)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    table.to_pandas().to_parquet(out, index=False)
+    return out
+
+
+def fetch_universe(universe: str, dt: str = 'D') -> tuple[Path, int]:
+    """bundle 落盘 + bundle 内全部 single 切片.
+
+    返 (bundle_path, n_single_written). 单次读入 bundle 后 groupby 切片,
+    避免逐股 pyarrow 全文件扫描 (5000 股 × 422MB = 不可接受).
+    """
+    bundle_path = download_bundle(universe, dt)
+    import pyarrow.parquet as pq
+    df = pq.read_table(bundle_path, columns=list(_SINGLE_COLUMNS)).to_pandas()
+    n_ok = 0
+    for code, sub in df.groupby('stockCode'):
+        out = _cache._dataset_path(str(code), 'D', 1)
+        out.parent.mkdir(parents=True, exist_ok=True)
+        sub.to_parquet(out, index=False)
+        n_ok += 1
+    return bundle_path, n_ok
+
+
+# ---- CLI ----
+
+def _cmd_list(_args) -> int:
+    try:
+        items = list_available()
+    except RuntimeError as e:
+        print(f'[ERR] {e}', file=sys.stderr)
+        return 2
+    if not items:
+        print('[OK] server 内置预构建 = 空 (或 builtin/enabled 全关)', file=sys.stderr)
+        return 0
+    print(f'[OK] server 内置预构建 ({len(items)} 条):', file=sys.stderr)
+    for d in items:
+        print(f'  {d.get("universe_code"):20s} {d.get("data_type_code"):12s} '
+              f'builtin={d.get("builtin")} enabled={d.get("enabled")}',
+              file=sys.stderr)
+    return 0
+
+
+def _cmd_fetch(args) -> int:
+    target = args.target
+    try:
+        if target == 'bundle':
+            p = download_bundle(args.universe, dt=args.dt)
+            size_mb = round(p.stat().st_size / 1e6, 1) if p.exists() else 0
+            print(f'[OK] bundle {args.universe}/{args.dt} → {p} ({size_mb} MB)',
+                  file=sys.stderr)
+        elif target == 'single':
+            p = download_single(args.symbol, period=args.period,
+                                start=args.start, end=args.end)
+            print(f'[OK] single {args.symbol} ({args.period}) → {p}', file=sys.stderr)
+        elif target == 'all':
+            bundle_path, n_ok = fetch_universe(args.universe, dt=args.dt)
+            print(f'[OK] all {args.universe}/{args.dt} → bundle={bundle_path}, '
+                  f'{n_ok} single 落盘', file=sys.stderr)
+        else:
+            print(f'[ERR] 未知 target={target!r}', file=sys.stderr)
+            return 2
+    except RuntimeError as e:
+        print(f'[ERR] {e}', file=sys.stderr)
+        return 2
+    return 0
+
+
+def main(argv: list[str] | None = None) -> int:
+    p = argparse.ArgumentParser(
+        prog='prebuilt_downloader',
+        description='下载 hamuna prebuilt 数据集 (v2 自有: server 加密凭据 + boto3 直连)',
+    )
+    sub = p.add_subparsers(dest='cmd', required=True)
+
+    sub.add_parser('list', help='列出 server 内置预构建数据集').set_defaults(func=_cmd_list)
+
+    f = sub.add_parser('fetch', help='下载 prebuilt (bundle / single / all)')
+    f.add_argument('target', choices=['bundle', 'single', 'all'])
+    # bundle + all 用 universe
+    f.add_argument('--universe', help='bundle/all 用: universe 名 (e.g. all_a / etf / convertible_bond)')
+    f.add_argument('--dt', default='D', help='bundle/all 用: 日期类型 (默认 D)')
+    # single 用 symbol
+    f.add_argument('--symbol', help='single 用: 标的代码 (e.g. 600000.SH)')
+    f.add_argument('--period', default='1d', help='single 用: 周期 (默认 1d)')
+    f.add_argument('--start', help='single 用: 起始 YYYYMMDD (bundle 切片全量, 不截断)')
+    f.add_argument('--end', help='single 用: 截止 YYYYMMDD (bundle 切片全量, 不截断)')
+    f.set_defaults(func=_cmd_fetch)
+
+    args = p.parse_args(argv)
+    return args.func(args)
+
+
+if __name__ == '__main__':
+    sys.exit(main())