Yuxuan Yan
79 lines
3.1 KiB
Python
79 lines
3.1 KiB
Python
"""CLI for data download phase."""
|
|
|
|
import click
|
|
from datetime import date
|
|
|
|
from pipeline.data.downloader import download_minute_universe, download_universe
|
|
|
|
|
|
@click.group(name="data")
|
|
def data():
|
|
"""Download and manage market data."""
|
|
|
|
|
|
@data.command("download")
|
|
@click.option(
|
|
"--universe", default="csi500",
|
|
help="Which universe: hs300, csi500, all (~5000 A-shares), file path, or comma-separated symbols",
|
|
)
|
|
@click.option("--start-date", default="2017-01-01", help="Start date YYYY-MM-DD")
|
|
@click.option("--end-date", default=str(date.today()), help="End date YYYY-MM-DD")
|
|
@click.option("--output-dir", default="data/daily_bars", help="Root for the partitioned dataset")
|
|
@click.option("--symbols", default=0, type=int, help="Max symbols (0=all)")
|
|
@click.option("--chunk-size", default=300, type=int, help="Symbols per durability flush")
|
|
@click.option("--adjust", default="qfq", help="Price adjust: qfq, hfq, or none")
|
|
def download(universe, start_date, end_date, output_dir, symbols, chunk_size, adjust):
|
|
"""Download daily bars into a month-partitioned parquet dataset.
|
|
|
|
Writes ``{output_dir}/{universe}/month=YYYY-MM/*.pq``. Point ``alpha
|
|
compute --data-path`` at that dataset directory.
|
|
"""
|
|
stats = download_universe(
|
|
universe=universe,
|
|
start_date=start_date,
|
|
end_date=end_date,
|
|
output_dir=output_dir,
|
|
max_symbols=symbols,
|
|
chunk_size=chunk_size,
|
|
adjust=adjust,
|
|
)
|
|
click.echo(
|
|
f"\nSummary: {stats['n_symbols']}/{stats['n_requested']} symbols, "
|
|
f"{stats['n_rows']:,} bars, {stats['date_min']} → {stats['date_max']}"
|
|
)
|
|
click.echo(f"Dataset: {stats['dataset_path']}")
|
|
|
|
|
|
@data.command("download-minute")
|
|
@click.option(
|
|
"--universe", default="csi500",
|
|
help="Which universe: hs300, csi500, all (~5000 A-shares), file path, or comma-separated symbols",
|
|
)
|
|
@click.option("--start-date", default="2017-01-01", help="Start date YYYY-MM-DD")
|
|
@click.option("--end-date", default=str(date.today()), help="End date YYYY-MM-DD")
|
|
@click.option("--output-dir", default="data/minute_bars", help="Root for the partitioned dataset")
|
|
@click.option("--symbols", default=0, type=int, help="Max symbols (0=all)")
|
|
@click.option("--chunk-size", default=100, type=int, help="Symbols per durability flush")
|
|
@click.option("--frequency", default="5", help="Minute frequency: 5, 15, 30, or 60")
|
|
def download_minute(universe, start_date, end_date, output_dir, symbols, chunk_size, frequency):
|
|
"""Download raw Baostock minute bars into a partitioned parquet dataset.
|
|
|
|
Writes ``{output_dir}/{universe}/frequency=5m/month=YYYY-MM/*.pq`` for the
|
|
default 5-minute frequency.
|
|
"""
|
|
stats = download_minute_universe(
|
|
universe=universe,
|
|
start_date=start_date,
|
|
end_date=end_date,
|
|
output_dir=output_dir,
|
|
max_symbols=symbols,
|
|
chunk_size=chunk_size,
|
|
frequency=frequency,
|
|
)
|
|
click.echo(
|
|
f"\nSummary: {stats['n_symbols']}/{stats['n_requested']} symbols, "
|
|
f"{stats['n_rows']:,} bars, {stats['date_min']} → {stats['date_max']}, "
|
|
f"frequency={stats['frequency']}"
|
|
)
|
|
click.echo(f"Dataset: {stats['dataset_path']}")
|