• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python3
2
3import argparse
4import ast
5import sys
6import os
7from time import time
8
9try:
10    import memory_profiler
11except ModuleNotFoundError:
12    print(
13        "Please run `make venv` to create a virtual environment and install"
14        " all the dependencies, before running this script."
15    )
16    sys.exit(1)
17
18sys.path.insert(0, os.getcwd())
19from scripts.test_parse_directory import parse_directory
20
21argparser = argparse.ArgumentParser(
22    prog="benchmark", description="Reproduce the various pegen benchmarks"
23)
24argparser.add_argument(
25    "--target",
26    action="store",
27    choices=["xxl", "stdlib"],
28    default="xxl",
29    help="Which target to use for the benchmark (default is xxl.py)",
30)
31
32subcommands = argparser.add_subparsers(title="Benchmarks", dest="subcommand")
33command_compile = subcommands.add_parser(
34    "compile", help="Benchmark parsing and compiling to bytecode"
35)
36command_parse = subcommands.add_parser("parse", help="Benchmark parsing and generating an ast.AST")
37
38
39def benchmark(func):
40    def wrapper(*args):
41        times = list()
42        for _ in range(3):
43            start = time()
44            result = func(*args)
45            end = time()
46            times.append(end - start)
47        memory = memory_profiler.memory_usage((func, args))
48        print(f"{func.__name__}")
49        print(f"\tTime: {sum(times)/3:.3f} seconds on an average of 3 runs")
50        print(f"\tMemory: {max(memory)} MiB on an average of 3 runs")
51        return result
52
53    return wrapper
54
55
56@benchmark
57def time_compile(source):
58    return compile(source, "<string>", "exec")
59
60
61@benchmark
62def time_parse(source):
63    return ast.parse(source)
64
65
66def run_benchmark_xxl(subcommand, source):
67    if subcommand == "compile":
68        time_compile(source)
69    elif subcommand == "parse":
70        time_parse(source)
71
72
73def run_benchmark_stdlib(subcommand):
74    modes = {"compile": 2, "parse": 1}
75    for _ in range(3):
76        parse_directory(
77            "../../Lib",
78            verbose=False,
79            excluded_files=["*/bad*", "*/lib2to3/tests/data/*",],
80            short=True,
81            mode=modes[subcommand],
82        )
83
84
85def main():
86    args = argparser.parse_args()
87    subcommand = args.subcommand
88    target = args.target
89
90    if subcommand is None:
91        argparser.error("A benchmark to run is required")
92
93    if target == "xxl":
94        with open(os.path.join("data", "xxl.py"), "r") as f:
95            source = f.read()
96            run_benchmark_xxl(subcommand, source)
97    elif target == "stdlib":
98        run_benchmark_stdlib(subcommand)
99
100
101if __name__ == "__main__":
102    main()
103