#!/usr/bin/env python3 from pathlib import Path from types import SimpleNamespace import click HERE = Path(__file__).resolve().parent def parse_alloc_data(alloc_data): parsed_data = {} for line in alloc_data: ident, allocs, calls = line.strip().split(" ") allocs = int(allocs) calls = int(calls) filename, lineno = ident.split(":") lineno = int(lineno) filedata = parsed_data.setdefault(filename, {}) filedata[lineno] = { "total_allocs": allocs, "total_calls": calls, "avg_allocs": allocs / calls, } return parsed_data @click.group() @click.pass_context @click.option("-a", "--alloc-data", type=click.File(), default="src/alloc_data.txt") @click.option("-t", "--type", type=click.Choice(("total", "avg")), default="avg") def cli(ctx, alloc_data, type): ctx.obj = SimpleNamespace(data=parse_alloc_data(alloc_data), type=type) def _normalize_filename(filename): if filename.startswith("src/"): return filename[4:] return filename @cli.command() @click.pass_obj @click.argument("filename") def annotate(obj, filename): filename = _normalize_filename(filename) if obj.type == "total": alloc_str = lambda line: str(line["total_allocs"]) else: alloc_str = lambda line: f"{line['avg_allocs']:.2f}" filedata = obj.data[filename] linedata = {lineno: alloc_str(line) for lineno, line in filedata.items()} maxlen = max(len(l) for l in linedata.values()) lineno = 0 for line in open("src/" + filename): lineno += 1 linecount = linedata.get(lineno, "") print(f"{linecount:>{maxlen}} {line}", end="") def _list(obj, sort_by="avg_allocs", reverse=False): return sorted( ( ( filename, sum(line["avg_allocs"] for line in lines.values()), sum(line["total_allocs"] for line in lines.values()), ) for filename, lines in obj.data.items() ), key=lambda x: x[1 if sort_by == "avg_allocs" else 2], reverse=reverse, ) @cli.command() @click.pass_obj @click.option("-r", "--reverse", is_flag=True) def list(obj, reverse): if obj.type == "total": field = "total_allocs" format_num = lambda l: f"{l[2]}" else: field = "avg_allocs" format_num = lambda l: f"{l[1]:.2f}" file_sums = _list(obj, field, reverse) maxlen = max(len(format_num(l)) for l in file_sums) for l in file_sums: num_str = format_num(l) filename = l[0] print(f"{num_str:>{maxlen}} {filename}") class HtmlTable: def __init__(self, f): self.f = f def __enter__(self): self.f.write("<table>") return self def __exit__(self, type, value, traceback): self.f.write("</table>") def tr(self, *tds): self.f.write("<tr>") for td in tds: if isinstance(td, tuple): self.f.write(f"<td {td[0]}><tt>{td[1]}</tt></td>") else: self.f.write(f"<td><tt>{td}</tt></td>") self.f.write("</tr>") @cli.command() @click.pass_obj @click.argument("htmldir") def html(obj, htmldir): file_sums = _list(obj, "total_allocs", reverse=True) style_grey = "style='color: grey'" style_right = "style='text-align: right'" with open(f"{htmldir}/index.html", "w") as f: f.write("<html>") f.write( f"<h3>Total allocations: {sum(total_sum for _, _, total_sum in file_sums)}</h3>" ) with HtmlTable(f) as table: table.tr((style_right, "avg"), (style_right, "total"), "") for filename, avg_sum, total_sum in file_sums: table.tr( (style_right, f"{avg_sum:.2f}"), (style_right, total_sum), f"<a href='{filename}.html'>{filename}</a>", ) f.write("</html>") for filename in file_sums: filename = _normalize_filename(filename[0]) htmlfile = Path(htmldir) / filename htmlfile.parent.mkdir(parents=True, exist_ok=True) with open(str(htmlfile) + ".html", "w") as f: filedata = obj.data[filename] f.write(f"<html><title>{filename}</title>") with HtmlTable(f) as table: table.tr( (style_grey, "#"), (style_right, "avg"), (style_right, "total"), "" ) lineno = 0 for line in open(HERE.parent / "src" / filename): line = line.rstrip("\n").replace(" ", " ") lineno += 1 total = filedata.get(lineno, {}).get("total_allocs", 0) avg = filedata.get(lineno, {}).get("avg_allocs", 0) table.tr( (style_grey, lineno), (style_right, f"{avg:.2f}"), (style_right, total), line, ) f.write("</html>") if __name__ == "__main__": cli()