Add heatmap_visualization/src/heatmap_visualization/cli.py

heatmap_visualization/src/heatmap_visualization/cli.py

@@ -0,0 +1,70 @@
+import argparse
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+from heatmap_visualization import core
+
+# Configure logging for CI and debugging
+logging.basicConfig(level=logging.INFO, format='[%(levelname)s] %(message)s')
+logger = logging.getLogger(__name__)
+
+
+def _validate_input_data(data: Any) -> pd.DataFrame:
+    """Validate that the input data conforms to the expected LogData structure."""
+    if isinstance(data, list):
+        try:
+            df = pd.DataFrame(data)
+        except Exception as e:
+            raise ValueError(f"Could not convert input list to DataFrame: {e}")
+    elif isinstance(data, pd.DataFrame):
+        df = data.copy()
+    else:
+        raise TypeError("Input data must be a list of dicts or a pandas.DataFrame.")
+
+    required_fields = {"worker_start_offset", "expires_at_dist_hours", "retry_total_overhead_ms"}
+    missing = required_fields - set(df.columns)
+    if missing:
+        raise ValueError(f"Input data missing required fields: {', '.join(missing)}")
+
+    # Validate types
+    for field in required_fields:
+        if not pd.api.types.is_numeric_dtype(df[field]):
+            raise ValueError(f"Field '{field}' must be numeric.")
+    return df
+
+
+def main() -> None:
+    """Command-line entry point for generating the Run 26 heatmap visualization."""
+    parser = argparse.ArgumentParser(description="Generate a latency heatmap from run log data.")
+    parser.add_argument("--input", required=True, help="Path to input log JSON file.")
+    parser.add_argument("--output", required=True, help="Path to output PNG file.")
+    args = parser.parse_args()
+
+    input_path = Path(args.input)
+    output_path = Path(args.output)
+
+    if not input_path.exists():
+        logger.error(f"Input file not found: {input_path}")
+        raise FileNotFoundError(f"Input file not found: {input_path}")
+
+    logger.info(f"Loading JSON data from {input_path}")
+    with input_path.open("r", encoding="utf-8") as f:
+        data = json.load(f)
+
+    df = _validate_input_data(data)
+
+    logger.info("Generating heatmap figure...")
+    fig = core.generate_heatmap(df)
+
+    logger.info(f"Saving heatmap to {output_path}")
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(output_path, dpi=300, bbox_inches="tight")
+
+    logger.info("Heatmap generation completed successfully.")
+
+
+if __name__ == "__main__":
+    main()