"""Inventory a local Tushare mirror and build/verify a frozen Qlib provider.""" from __future__ import annotations import argparse import json from pathlib import Path import sys from typing import Any, Mapping, Optional from adapters.tushare_local import ( TushareMirrorError, build_qlib_provider, inventory_mirror, verify_qlib_provider, ) def _write_json(path: Optional[str], payload: Mapping[str, Any]) -> None: if path is None: return destination = Path(path).expanduser().resolve() if destination.exists() and destination.is_dir(): raise ValueError(f"output-json points to a directory: {destination}") destination.parent.mkdir(parents=True, exist_ok=True) destination.write_text( json.dumps( dict(payload), ensure_ascii=False, indent=2, sort_keys=True, allow_nan=False, ) + "\n", encoding="utf-8", ) def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=__doc__) subparsers = parser.add_subparsers(dest="command", required=True) inventory = subparsers.add_parser("inventory") inventory.add_argument("--mirror-root", required=True) inventory.add_argument("--skip-file-verification", action="store_true") inventory.add_argument("--output-json") build = subparsers.add_parser("build") build.add_argument("--mirror-root", required=True) build.add_argument("--output-dir", required=True) build.add_argument("--start", required=True) build.add_argument("--end", required=True) build.add_argument("--market-name", default="tushare_a") build.add_argument("--benchmark-symbol", default="SH999999") build.add_argument("--minimum-observations", type=int, default=60) build.add_argument( "--allow-unadjusted", action="store_true", help="required while adj_factor is absent; never grants investment credit", ) build.add_argument( "--ohlc-policy", choices=("fail", "expand-range"), default="fail", ) build.add_argument("--output-json") verify = subparsers.add_parser("verify") verify.add_argument("provider_dir") verify.add_argument("--output-json") return parser def _main(argv: Optional[list[str]] = None) -> int: parser = build_parser() args = parser.parse_args(argv) try: if args.command == "inventory": result = inventory_mirror( args.mirror_root, verify_files=not args.skip_file_verification, ) elif args.command == "build": result = build_qlib_provider( args.mirror_root, args.output_dir, start=args.start, end=args.end, market_name=args.market_name, benchmark_symbol=args.benchmark_symbol, minimum_observations=args.minimum_observations, allow_unadjusted=args.allow_unadjusted, expand_range=args.ohlc_policy == "expand-range", ) else: result = verify_qlib_provider(args.provider_dir) _write_json(args.output_json, result) except (OSError, ValueError, TushareMirrorError) as exc: parser.error(str(exc)) print( json.dumps( result, ensure_ascii=False, indent=2, sort_keys=True, allow_nan=False, ) ) return 0 if __name__ == "__main__": sys.exit(_main())