chinese-equity-quant/pipeline/derived/cli.py

"""CLI for daily derived-data ingestion and computation."""

import click
import pandas as pd

from pipeline.derived.compute import (
    DERIVED_KEY_COLUMNS,
    compute_derived,
    read_derived_frame,
    write_derived_frame,
)
from pipeline.derived.registry import (
    available_derived,
    load_derived_module,
)


@click.group(name="derived")
def derived():
    """Ingest, compute, and validate daily derived data."""


def _coerce(value: str):
    """Best-effort coercion of a CLI string to int, then float, else str."""
    for cast in (int, float):
        try:
            return cast(value)
        except ValueError:
            continue
    return value


def _parse_params(pairs: tuple[str, ...]) -> dict:
    """Parse repeated ``name=value`` options into a params dict."""
    params: dict = {}
    for pair in pairs:
        if "=" not in pair:
            raise click.BadParameter(f"--param must be name=value, got '{pair}'")
        key, value = pair.split("=", 1)
        params[key.strip()] = _coerce(value.strip())
    return params


def _read_optional_parquet(path: str | None) -> pd.DataFrame | None:
    return None if path is None else pd.read_parquet(path)


def _summarize(result: pd.DataFrame) -> str:
    value_cols = [col for col in result.columns if col not in DERIVED_KEY_COLUMNS]
    return f"{len(result):,} rows, {len(value_cols)} columns"


@derived.command("list")
@click.option(
    "--derived-module", "derived_modules", multiple=True,
    help="External module(s) to import first (dotted path or .py file)",
)
def list_(derived_modules):
    """List the registered derived-data plugin types."""
    for spec in derived_modules:
        load_derived_module(spec)
    for name in available_derived():
        click.echo(name)


@derived.command("validate")
@click.option("--input-path", required=True, help="CSV/parquet file or parquet dataset to validate")
def validate(input_path):
    """Validate a daily derived-data file without writing output."""
    try:
        result = read_derived_frame(input_path)
    except Exception as exc:
        raise click.ClickException(str(exc)) from exc
    click.echo(f"Valid derived data: {input_path}  ({_summarize(result)})")


@derived.command("ingest")
@click.option("--input-path", required=True, help="CSV/parquet file to ingest")
@click.option("--derived-name", required=True, help="Name for this derived-data output file")
@click.option("--output-dir", default="derived", help="Directory to save derived parquet")
def ingest(input_path, derived_name, output_dir):
    """Ingest a user-provided daily derived-data CSV/parquet file."""
    try:
        result = read_derived_frame(input_path)
        out_path = write_derived_frame(result, derived_name, output_dir=output_dir)
    except Exception as exc:
        raise click.ClickException(str(exc)) from exc
    click.echo(f"Saved derived data: {out_path}  ({_summarize(result)})")


@derived.command("compute")
@click.option("--daily-path", default=None, help="Optional daily data parquet/dataset")
@click.option("--minute-path", default=None, help="Optional minute parquet/dataset")
@click.option("--derived-type", required=True, help="Registry key of the derived-data plugin")
@click.option("--derived-name", required=True, help="Name for this derived-data output file")
@click.option("--output-dir", default="derived", help="Directory to save derived parquet")
@click.option(
    "--derived-module", "derived_modules", multiple=True,
    help="External module(s) to import so their derived-data plugins register",
)
@click.option(
    "--param", "extra_params", multiple=True,
    help="Extra derived-data constructor param as name=value (repeatable)",
)
def compute(
    daily_path,
    minute_path,
    derived_type,
    derived_name,
    output_dir,
    derived_modules,
    extra_params,
):
    """Compute one daily derived-data file from daily and/or minute inputs."""
    for spec in derived_modules:
        load_derived_module(spec)

    options = available_derived()
    if derived_type not in options:
        raise click.BadParameter(
            f"Unknown derived-type '{derived_type}'. Available: {options}. "
            f"Use --derived-module to register an external derived-data plugin.",
            param_hint="--derived-type",
        )
    if daily_path is None and minute_path is None:
        raise click.UsageError("At least one of --daily-path or --minute-path is required")

    daily = _read_optional_parquet(daily_path)
    if daily_path:
        click.echo(f"Loaded daily data: {len(daily):,} rows from {daily_path}")
    minute = _read_optional_parquet(minute_path)
    if minute_path:
        click.echo(f"Loaded minute bars: {len(minute):,} rows from {minute_path}")

    try:
        result = compute_derived(
            derived_type=derived_type,
            daily=daily,
            minute=minute,
            **_parse_params(extra_params),
        )
        out_path = write_derived_frame(result, derived_name, output_dir=output_dir)
    except Exception as exc:
        raise click.ClickException(str(exc)) from exc
    click.echo(f"Saved derived data: {out_path}  ({_summarize(result)})")