plotting.py

import json
import matplotlib.pyplot as plt

def plot_benchmark_results(
    json_file: str,
    labels: list[str],
    output_file: str | None = None,
    title: str | None = None,
) -> None:
    """
    Plots benchmark results from a JSON file generated by hyperfine.

    Args:
        json_file: Path to the JSON file with benchmark results.
        labels: List of labels for the plot legend.
        output_file: Optional path to save the generated image. If None, the plot is shown instead.
        title: Optional title for the plot.
    """
    with open(json_file, encoding="utf-8") as f:
        results = json.load(f)["results"]

    times = [b["times"] for b in results]

    plt.figure(figsize=(10, 6), constrained_layout=True)
    boxplot = plt.boxplot(times, vert=True, patch_artist=True)
    cmap = plt.get_cmap("rainbow")
    colors = [cmap(val / len(times)) for val in range(len(times))]

    for patch, color in zip(boxplot["boxes"], colors):
        patch.set_facecolor(color)

    if title:
        plt.title(title)

    plt.legend(handles=boxplot["boxes"], labels=labels, loc="best", fontsize="medium")
    plt.ylabel("Time [s]")
    plt.ylim(0, None)
    plt.xticks(list(range(1, len(labels) + 1)), labels, rotation=45)
    
    if output_file:
        plt.savefig(output_file)
    else:
        plt.show()
    plt.close()