"""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 PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from adapters.tushare_local import ( TushareMirrorError, build_qlib_provider, inventory_mirror, verify_qlib_provider, ) def _assert_json_outside_provider( output_json: Optional[str], provider_dir: str, ) -> None: if output_json is None: return destination = Path(output_json).expanduser().resolve() provider = Path(provider_dir).expanduser().resolve() if destination == provider or provider in destination.parents: raise ValueError( "output-json must be outside the immutable provider directory" ) 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( "--benchmark-index-code", help="use real Tushare index_daily data, for example 000905.SH", ) build.add_argument( "--universe-index-code", help="build point-in-time market intervals from index_weight snapshots", ) build.add_argument( "--expected-constituent-count", type=int, help=( "required for unknown index codes; known indices such as " "000905.SH use a built-in exact count" ), ) 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": _assert_json_outside_provider( args.output_json, args.output_dir, ) 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", benchmark_index_code=args.benchmark_index_code, universe_index_code=args.universe_index_code, expected_constituent_count=args.expected_constituent_count, ) else: _assert_json_outside_provider( args.output_json, args.provider_dir, ) 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())