1  #!/usr/bin/env python3
 2  
 3  import argparse
 4  import sys
 5  import re
 6  from pathlib import Path
 7  from typing import Optional
 8  from collections.abc import Callable, Iterable
 9  from tabulate import tabulate
10  from dataclasses import asdict
11  
12  topdir = Path(__file__).parent.parent
13  sys.path.insert(0, str(topdir))
14  
15  
16  from transactron.profiler import Profile, RunStat, RunStatNode  # noqa: E402
17  
18  
19  def process_stat_tree(
20      xs: Iterable[RunStatNode], recursive: bool, ret: Optional[list[tuple]] = None, depth=0
21  ) -> list[tuple]:
22      if ret is None:
23          ret = list[tuple]()
24      for x in xs:
25          row = asdict(x.stat)
26          if recursive and depth:
27              row["name"] = (2 * depth - 1) * "-" + " " + row["name"]
28          ret.append(tuple(row.values()))
29          if recursive and x.callers:
30              process_stat_tree(x.callers.values(), recursive, ret, depth + 1)
31      return ret
32  
33  
34  def filter_nodes(nodes: list[RunStatNode], key: Callable[[RunStat], str], regex: str):
35      pattern = re.compile(regex)
36      return [node for node in nodes if pattern.search(key(node.stat))]
37  
38  
39  def sort_node(node: RunStatNode, sort_order: str):
40      node.callers = dict(sorted(node.callers.items(), key=lambda node: asdict(node[1].stat)[sort_order]))
41      for node2 in node.callers.values():
42          sort_node(node2, sort_order)
43  
44  
45  def main():
46      parser = argparse.ArgumentParser()
47      parser.add_argument("-g", "--call-graph", action="store_true", help="Show call graph")
48      parser.add_argument("-s", "--sort", choices=["name", "locked", "run"], default="name", help="Sort by column")
49      parser.add_argument(
50          "-m", "--mode", choices=["transactions", "methods"], default="transactions", help="Profile display mode"
51      )
52      parser.add_argument("-f", "--filter-name", help="Filter by name, regular expressions can be used")
53      parser.add_argument("-l", "--filter-loc", help="Filter by source location, regular expressions can be used")
54      parser.add_argument("file_name", nargs=1)
55  
56      args = parser.parse_args()
57  
58      profile = Profile.decode(args.file_name[0])
59  
60      recursive = args.call_graph
61  
62      if args.mode == "transactions":
63          nodes = profile.analyze_transactions(recursive=recursive)
64      elif args.mode == "methods":
65          nodes = profile.analyze_methods(recursive=recursive)
66      else:
67          assert False
68  
69      headers = ["name", "source location", "locked", "run"]
70  
71      nodes.sort(key=lambda node: asdict(node.stat)[args.sort])
72      for node in nodes:
73          sort_node(node, args.sort)
74  
75      if args.filter_name:
76          nodes = filter_nodes(nodes, lambda stat: stat.name, args.filter_name)
77      if args.filter_loc:
78          nodes = filter_nodes(nodes, lambda stat: stat.src_loc, args.filter_loc)
79  
80      print(tabulate(process_stat_tree(nodes, recursive), headers=headers))
81  
82  
83  if __name__ == "__main__":
84      main()