quickbench.py

#!/usr/bin/env python3
"""
quickbench -- Benchmark shell commands. Like hyperfine, but Python + zero deps.

Runs commands multiple times, reports statistical analysis. Compare commands
side-by-side. Cross-platform (Windows, macOS, Linux). No pip installs needed.

Usage:
    py quickbench.py "python -c 'sum(range(10000))'"                   # Benchmark one command
    py quickbench.py "sort big.txt" "sort -S 50% big.txt" -n 20        # Compare two commands
    py quickbench.py "ls -R" "find . -type f" -w 3                     # 3 warmup runs
    py quickbench.py "make build" --json                                # JSON output
    py quickbench.py "cmd1" "cmd2" "cmd3" -n 50 --export results.csv   # CSV export
"""

import argparse
import json
import math
import os
import subprocess
import sys
import time

# Colors
RESET = "\033[0m"
BOLD = "\033[1m"
DIM = "\033[2m"
GREEN = "\033[32m"
YELLOW = "\033[33m"
CYAN = "\033[36m"
MAGENTA = "\033[35m"
RED = "\033[31m"


def color_supported() -> bool:
    if os.environ.get("NO_COLOR"):
        return False
    if sys.platform == "win32":
        return bool(os.environ.get("TERM") or os.environ.get("WT_SESSION"))
    return hasattr(sys.stdout, "isatty") and sys.stdout.isatty()


USE_COLOR = color_supported()


def c(code: str, text: str) -> str:
    return f"{code}{text}{RESET}" if USE_COLOR else text


def run_command(cmd: str, shell_path: str | None = None) -> tuple[float, int]:
    """Run a command and return (elapsed_seconds, exit_code)."""
    start = time.perf_counter()
    try:
        if sys.platform == "win32":
            result = subprocess.run(cmd, shell=True, capture_output=True)
        else:
            result = subprocess.run(
                cmd, shell=True, capture_output=True,
                executable=shell_path or "/bin/sh"
            )
        elapsed = time.perf_counter() - start
        return elapsed, result.returncode
    except Exception:
        elapsed = time.perf_counter() - start
        return elapsed, 1


def compute_stats(times: list[float]) -> dict:
    """Compute statistics for a list of timing measurements."""
    n = len(times)
    if n == 0:
        return {"mean": 0, "median": 0, "stddev": 0, "min": 0, "max": 0, "n": 0}

    sorted_t = sorted(times)
    mean = sum(times) / n

    if n % 2 == 0:
        median = (sorted_t[n // 2 - 1] + sorted_t[n // 2]) / 2
    else:
        median = sorted_t[n // 2]

    if n > 1:
        variance = sum((t - mean) ** 2 for t in times) / (n - 1)
        stddev = math.sqrt(variance)
    else:
        stddev = 0.0

    return {
        "mean": mean,
        "median": median,
        "stddev": stddev,
        "min": sorted_t[0],
        "max": sorted_t[-1],
        "n": n,
    }


def format_time(seconds: float) -> str:
    """Format time in human-readable units."""
    if seconds < 0.001:
        return f"{seconds * 1_000_000:.1f} us"
    elif seconds < 1.0:
        return f"{seconds * 1000:.1f} ms"
    elif seconds < 60:
        return f"{seconds:.3f} s"
    else:
        m = int(seconds // 60)
        s = seconds % 60
        return f"{m}m {s:.1f}s"


def format_time_aligned(seconds: float) -> str:
    """Format with consistent width for table alignment."""
    if seconds < 0.001:
        return f"{seconds * 1_000_000:>8.1f} us"
    elif seconds < 1.0:
        return f"{seconds * 1000:>8.1f} ms"
    elif seconds < 60:
        return f"{seconds:>8.3f} s "
    else:
        return f"{seconds:>8.1f} s "


def progress_bar(current: int, total: int, width: int = 20) -> str:
    filled = int(width * current / total) if total > 0 else 0
    bar = "#" * filled + "-" * (width - filled)
    return f"[{bar}]"


def benchmark_command(cmd: str, runs: int, warmup: int, label: str,
                      show_progress: bool = True) -> dict:
    """Benchmark a single command."""
    # Warmup
    if warmup > 0 and show_progress:
        print(c(DIM, f"  Warming up {label}... ({warmup} runs)"), end="", flush=True)
        for _ in range(warmup):
            run_command(cmd)
        print(c(DIM, " done"))

    times = []
    failures = 0

    for i in range(runs):
        if show_progress:
            pbar = progress_bar(i + 1, runs)
            pct = (i + 1) / runs * 100
            print(f"\r  {c(CYAN, label)} {pbar} {pct:5.1f}%", end="", flush=True)

        elapsed, code = run_command(cmd)
        times.append(elapsed)
        if code != 0:
            failures += 1

    if show_progress:
        print(f"\r  {c(CYAN, label)} {progress_bar(runs, runs)} {c(GREEN, 'done')}   ")

    stats = compute_stats(times)
    stats["command"] = cmd
    stats["label"] = label
    stats["failures"] = failures
    stats["times"] = times
    return stats


def print_results(all_stats: list[dict]):
    """Print benchmark results in a formatted table."""
    print()
    print(c(BOLD, "  Benchmark Results"))
    print(c(DIM, "  " + "-" * 60))

    for i, stats in enumerate(all_stats):
        label = stats["label"]
        color = [CYAN, MAGENTA, YELLOW, GREEN][i % 4]

        print(f"\n  {c(BOLD + color, label)}")
        print(f"    Command: {c(DIM, stats['command'])}")
        print(f"    Mean:    {c(BOLD, format_time(stats['mean']))} +/- {format_time(stats['stddev'])}")
        print(f"    Median:  {format_time(stats['median'])}")
        print(f"    Range:   {format_time(stats['min'])} ... {format_time(stats['max'])}")
        print(f"    Runs:    {stats['n']}", end="")
        if stats["failures"] > 0:
            print(f"  ({c(RED, str(stats['failures']) + ' failed')})", end="")
        print()

    # Comparison
    if len(all_stats) > 1:
        print()
        print(c(BOLD, "  Comparison"))
        print(c(DIM, "  " + "-" * 60))

        fastest_idx = min(range(len(all_stats)), key=lambda i: all_stats[i]["mean"])
        fastest = all_stats[fastest_idx]

        for i, stats in enumerate(all_stats):
            color = [CYAN, MAGENTA, YELLOW, GREEN][i % 4]
            ratio = stats["mean"] / fastest["mean"] if fastest["mean"] > 0 else 0

            if i == fastest_idx:
                print(f"  {c(color, stats['label']):>30}  {format_time_aligned(stats['mean'])}  {c(GREEN, '(fastest)')}")
            else:
                print(f"  {c(color, stats['label']):>30}  {format_time_aligned(stats['mean'])}  {c(YELLOW, f'{ratio:.2f}x slower')}")

        print()
        winner = all_stats[fastest_idx]
        print(f"  {c(BOLD + GREEN, 'Winner:')} {winner['label']} ({format_time(winner['mean'])})")

    print()


def export_csv(all_stats: list[dict], path: str):
    """Export results to CSV."""
    with open(path, "w", encoding="utf-8") as f:
        f.write("label,command,mean_s,median_s,stddev_s,min_s,max_s,runs,failures\n")
        for s in all_stats:
            cmd_escaped = s["command"].replace('"', '""')
            f.write(f'"{s["label"]}","{cmd_escaped}",{s["mean"]:.6f},{s["median"]:.6f},'
                    f'{s["stddev"]:.6f},{s["min"]:.6f},{s["max"]:.6f},{s["n"]},{s["failures"]}\n')
    print(c(DIM, f"  Exported to {path}"))


def export_json(all_stats: list[dict]):
    """Print results as JSON."""
    output = []
    for s in all_stats:
        output.append({
            "label": s["label"],
            "command": s["command"],
            "mean": round(s["mean"], 6),
            "median": round(s["median"], 6),
            "stddev": round(s["stddev"], 6),
            "min": round(s["min"], 6),
            "max": round(s["max"], 6),
            "runs": s["n"],
            "failures": s["failures"],
            "times": [round(t, 6) for t in s["times"]],
        })
    print(json.dumps(output, indent=2))


def main():
    parser = argparse.ArgumentParser(
        description="quickbench -- benchmark shell commands",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("commands", nargs="+", help="Command(s) to benchmark")
    parser.add_argument("-n", "--runs", type=int, default=10, help="Number of runs (default: 10)")
    parser.add_argument("-w", "--warmup", type=int, default=1, help="Warmup runs (default: 1)")
    parser.add_argument("--json", action="store_true", help="Output as JSON")
    parser.add_argument("--export", metavar="FILE", help="Export results to CSV file")
    parser.add_argument("-l", "--labels", nargs="+", help="Custom labels for commands")
    parser.add_argument("-q", "--quiet", action="store_true", help="Minimal output")

    args = parser.parse_args()

    # Generate labels
    labels = []
    if args.labels and len(args.labels) >= len(args.commands):
        labels = args.labels[:len(args.commands)]
    else:
        for i, cmd in enumerate(args.commands):
            short = cmd[:40] + "..." if len(cmd) > 40 else cmd
            labels.append(short if len(args.commands) > 1 else "Command")

    if not args.json and not args.quiet:
        print()
        print(c(BOLD, "  quickbench"))
        print(c(DIM, f"  Benchmarking {len(args.commands)} command(s), {args.runs} runs each"))
        if args.warmup:
            print(c(DIM, f"  Warmup: {args.warmup} run(s)"))
        print()

    # Run benchmarks
    all_stats = []
    for i, (cmd, label) in enumerate(zip(args.commands, labels)):
        stats = benchmark_command(
            cmd, args.runs, args.warmup, label,
            show_progress=not args.json and not args.quiet
        )
        all_stats.append(stats)

    # Output
    if args.json:
        export_json(all_stats)
    else:
        print_results(all_stats)

    if args.export:
        export_csv(all_stats, args.export)


if __name__ == "__main__":
    main()