sortprecompile.py

#!/usr/bin/env python3
"""Sort a precompile_*.jl file by the `#= NNN.N ms =#` timing comment.

Lines look like:

    #=  465.4 ms =# precompile(Tuple{...})

Some lines are wrapped in ANSI color escapes (e.g. \\033[33m ... \\033[0m) and
end with `# recompile`; those escapes are stripped from all output.

Lines without a timing comment (headers, footers, blank lines) are kept and
emitted after the sorted block, in their original order.

Usage:
    sortprecompile [FILE]        # sort FILE in place, slowest first
    sortprecompile               # read stdin, write stdout
    sortprecompile -a FILE       # ascending (fastest first)
    sortprecompile -o OUT FILE   # write to OUT instead of in place
    sortprecompile -s FILE       # `# recompile` lines in their own section
    sortprecompile -m 10 FILE    # drop everything faster than 10 ms
"""
import argparse
import re
import sys

# ANSI/color escape sequences that show up as garbage (e.g. \033[33m ... \033[0m)
ANSI = re.compile(r"\x1b\[[0-9;]*m")

# match `#= <number> ms =#`
TIMING = re.compile(r"#=\s*([0-9]+(?:\.[0-9]+)?)\s*ms\s*=#")

# a recompile line ends with `# recompile`
RECOMPILE = re.compile(r"# recompile\s*$")

# a pre-existing section header from an earlier run; dropped so we own it
HEADER = re.compile(r"^#\s*Recompilation\s*$")


def main():
    ap = argparse.ArgumentParser(description="Sort precompile file by ms timing.")
    ap.add_argument("file", nargs="?", help="file to sort (default: stdin->stdout)")
    ap.add_argument("-a", "--ascending", action="store_true",
                    help="fastest first (default: slowest first)")
    ap.add_argument("-o", "--output", help="write result here instead of in place")
    ap.add_argument("-m", "--min-ms", type=float, default=0.0, metavar="MS",
                    help="drop any precompile line faster than MS milliseconds")
    ap.add_argument("-s", "--split-recompile", action="store_true",
                    help="put `# recompile` lines in a separate sorted section "
                         "under a `# Recompilation` header")
    args = ap.parse_args()

    if args.file:
        with open(args.file, encoding="utf-8") as f:
            lines = f.read().splitlines()
    else:
        lines = sys.stdin.read().splitlines()

    # strip ANSI color escapes (the garbage) from every line, and drop any
    # `# Recompilation` header left over from a previous run (we re-add it)
    lines = [ANSI.sub("", line) for line in lines]
    lines = [line for line in lines if not HEADER.match(line)]

    # everything after the last timed line is the footer (JuliaHub notes etc.),
    # preserved verbatim; leading blanks are dropped so re-runs don't grow them.
    timed_idx = [i for i, line in enumerate(lines) if TIMING.search(line)]
    last_t = timed_idx[-1] if timed_idx else -1
    footer = lines[last_t + 1:]
    while footer and not footer[0].strip():
        footer.pop(0)

    # entries live before the footer; drop stray blanks / leftover headers there
    timed = []      # (ms, original_index, line)
    recompile = []  # (ms, original_index, line) -- only when --split-recompile
    for i, line in enumerate(lines[:last_t + 1]):
        m = TIMING.search(line)
        if not m:
            continue
        ms = float(m.group(1))
        if ms < args.min_ms:
            continue
        bucket = recompile if args.split_recompile and RECOMPILE.search(line) else timed
        bucket.append((ms, i, line))

    # sort by ms with original index as a stable tie-break; negate for descending
    # so the tie-break is NOT reversed (that would make re-runs oscillate).
    sign = 1 if args.ascending else -1
    key = lambda t: (sign * t[0], t[1])
    timed.sort(key=key)
    recompile.sort(key=key)

    out = [line for _, _, line in timed]
    if recompile:
        out.append("")
        out.append("# Recompilation")
        out.extend(line for _, _, line in recompile)
    if footer:
        out.append("")
        out.extend(footer)
    text = "\n".join(out) + "\n"

    if args.output:
        with open(args.output, "w", encoding="utf-8") as f:
            f.write(text)
    elif args.file:
        with open(args.file, "w", encoding="utf-8") as f:
            f.write(text)
    else:
        sys.stdout.write(text)


if __name__ == "__main__":
    main()