generate_stats.py

#!/usr/bin/env -S uv run --script
#
# /// script
# requires-python = ">=3.13"
# dependencies = [
#   "tabulate",
# ]
# ///

import json
from collections import defaultdict
from pathlib import Path
from typing import NamedTuple

from tabulate import tabulate

EMOJI_SCALE = [
    ("❤️‍🩹", -25),
    ("🟥", -15),
    ("🔴", -5),
    ("⚪", 5),
    ("🟢", 15),
    ("🟩", 25),
    ("💚", None),
]


def number_machine(value: str) -> int:
    return (
        int(value[:-1])
        * {
            "k": 1_000,
            "M": 1_000_000,
            "G": 1_000_000_000,
        }[value[1]]
    )


class PythonVersion(NamedTuple):
    name: str
    version: str
    freethreaded: bool

    def __str__(self) -> str:
        version = f"{self.name} {self.version}"
        if self.freethreaded:
            version += "t"
        return version

    @property
    def order(self) -> object:
        return (
            ["CPython", "PyPy"].index(self.name),
            -int(self.version.split(".")[0]),
            -int(self.version.split(".")[1]),
            self.freethreaded,
        )


class Implementation(NamedTuple):
    name: str
    version: str


class Benchmark(NamedTuple):
    name: str
    size: str
    level: int
    python: PythonVersion
    implementation: Implementation
    iterations: int
    duration: float

    @property
    def order(self) -> object:
        return (
            *self.python.order,
            self.name,
            number_machine(self.size),
            self.level,
            self.implementation,
            self.duration / self.iterations,
        )


def read_benchmarks(path: Path) -> list[Benchmark]:
    benchmarks: list[Benchmark] = []
    with path.open() as f:
        for line in f:
            data = json.loads(line)
            benchmarks.append(
                Benchmark(
                    name=data["name"],
                    size=data["size"],
                    level=data["level"],
                    python=PythonVersion(
                        name=data["python"]["name"],
                        version=data["python"]["version"],
                        freethreaded=data["python"]["freethreaded"],
                    ),
                    implementation=Implementation(
                        name=data["implementation"]["name"],
                        version=data["implementation"]["version"],
                    ),
                    iterations=data["iterations"],
                    duration=data["duration"],
                )
            )
    return benchmarks


def run(path: Path) -> None:
    print("## Introduction")
    print()
    print("Comparing stdlib (or backports.zstd for Python before 3.14) against different libraries.")
    print()
    print(
        "Figures give timing comparison. For example, +42% means that the library needs 42% more time than stdlib/backports.zstd."
    )
    print()
    print("The reference time column indicates an average time for a single run.")
    print()
    print(
        "Emoji scale:",
        " ".join(h if v is None else f"{h} {v:+.0f}%" for h, v in EMOJI_SCALE),
    )
    print()
    benchmarks = read_benchmarks(path)
    versions = sorted({b.python for b in benchmarks}, key=lambda v: v.order)
    for version in versions:
        print("##", version)
        version_benchmarks = [b for b in benchmarks if b.python == version]
        print()
        implementations = {b.implementation for b in version_benchmarks}
        assert len(implementations) == len({i.name for i in implementations})
        has_backports_zstd = any(i.name == "backports.zstd" for i in implementations)
        has_stdin = any(i.name == "stdlib" for i in implementations)
        assert has_backports_zstd != has_stdin
        reference = "backports.zstd" if has_backports_zstd else "stdlib"
        print(f"Using as reference: `{reference}`")
        print()
        if len(implementations) < 2:
            print("*No other implementations on that Python version*")
            print()
            continue
        show_stats(
            sorted(version_benchmarks, key=lambda b: b.order),
            reference,
        )
        print()
        print(
            f"*{
                ', '.join(
                    sorted(
                        f'{i.name}=={i.version}'
                        for i in implementations
                        if i.name != 'stdlib'
                    )
                )
            }*"
        )
        print()


def format_value_cmp(diff: float) -> str:
    diff *= 100
    for h, threshold in EMOJI_SCALE:
        if threshold is None or diff < threshold:
            break
    return f"{h} {diff:=+6.2f}%"


def format_ref_time(timing: float, iterations: int) -> str:
    single_ns = timing * 1_000_000 / iterations
    if single_ns > 1_000_000:
        return f"{single_ns / 1_000_000:.2f}s"
    if single_ns > 1_000:
        return f"{single_ns / 1_000:.0f}ms"
    return "<1ms"


def show_stats(benchmarks: list[Benchmark], reference_implem: str) -> None:
    by_case = defaultdict(dict)
    implems = set()
    for b in benchmarks:
        by_case[b.name, b.size, b.level][b.implementation.name] = (
            b.duration,
            b.iterations,
        )
        implems.add(b.implementation.name)
    implems = sorted(implems)
    implems.remove(reference_implem)
    lines = []
    lines.append(["Case", "Ref time"] + implems)
    for case in sorted(by_case, key=lambda x: (x[0], number_machine(x[1]), x[2])):
        values = by_case[case]
        assert len({values[i][1] for i in implems}) == 1, implems
        values_cmp = {
            i: values[i][0] / values[reference_implem][0] - 1 for i in implems
        }
        lines.append(
            [
                f"{case[0]} {case[1]} level {case[2]}",
                format_ref_time(*values[reference_implem]),
            ]
            + [format_value_cmp(values_cmp[i]) for i in implems]
        )
    print(tabulate(lines, headers="firstrow", tablefmt="github"))


if __name__ == "__main__":
    from contextlib import redirect_stdout

    results_dir = Path(__file__).parent / "results"
    with (results_dir / "benchmark.md").open("w") as f, redirect_stdout(f):
        run(results_dir / "benchmark.ndjson")