helper script for tracing module imports with python -Ximporttime

importtrace.py

#!/usr/bin/env python

import argparse
import subprocess
import time
import sys

special_chars = '▀█░▒▓'

def imports_to_table(imports, title, module, runtime, width):
    print(title)
    table = []
    headers = ["name", "depth", "start(ms)", "self(ms)", "total(ms)", ""]
    for name, start, self_time, total_time, depth in sorted(imports, key=lambda x: (x[1], x[4])):
        s1 = int(start / runtime * width)
        s2 = int((start + total_time) / runtime * width) - s1
        s3 = width - s1 - s2
        if name in module:
            # highlight the named module(s)
            sc, se = '\033[94m', '\033[0m'
            char = special_chars[4]
        else:
            sc, se = '', ''
            char = special_chars[2]
        bar = (' ' * s1) + (sc + char * s2 + se) + (' ' * s3)
        indent = depth * "  "
        table.append([indent+name, depth, start, self_time, total_time, bar])
    try:
        import tabulate
        tabulate.PRESERVE_WHITESPACE = True
        print(tabulate.tabulate(table, headers))    
    except ImportError:
        print(f"{width*' '} | start(ms) self(ms) total(ms) depth name")
        print(f"{width*'-'} | ----------------------------------------")
        for row in table:
            # move the bar to the front to help with alignment
            indented_name, depth, start, self_time, total_time, bar = row
            metadata = [start, self_time, total_time, depth, indented_name]            
            metadata = f'{start:9d} {self_time:8d} {total_time:9d} {depth:5d} {indented_name}'
            print(f"{bar} | {metadata}")


def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("module", type=str, nargs="*", help="one or more module imports to trace")
    parser.add_argument("--max-depth", type=int, default=None, help="filter on max depth")
    parser.add_argument("--min-elapsed", type=int, default=10, help="filter on min elapsed")
    parser.add_argument("--percent", action="store_true", help="interpret --min-elapsed as a percent of total import time")
    parser.add_argument("--width", type=int, default=40, help="bar chart width")
    parser.add_argument("--exe", type=str, default=sys.executable, help="python exe")
    parser.add_argument("--script", type=str, default=None, help="interpret argument as a script instead of module name")
    parser.add_argument("--setup", type=str, default=None, help="prepend to cmd")
    args = parser.parse_args()

    # measure import time
    if args.script:
        if args.setup:
            shell = True
            cmd = args.setup
            cmd += f'; export PYTHONPROFILEIMPORTTIME=1; python {args.script}'
        else:
            shell = False
            cmd = [sys.executable, "-Ximporttime", args.script]
    else:
        imports = "; import ".join(args.module)
        if args.setup:
            shell = True
            cmd = args.setup
            cmd += f'; export PYTHONPROFILEIMPORTTIME=1; python -c "import {imports}"'
        else:
            shell = False
            cmd = [sys.executable, "-Ximporttime", "-c", f"import {imports}"]
    print(cmd)
    start = time.time()
    result = subprocess.run(cmd, shell=shell, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
    runtime = (time.time() - start) * 1000

    if args.percent:
        args.min_elapsed = runtime * args.min_elapsed / 100

    # process import time results
    imports = []
    start_stack = [0]

    LINE_PREFIX = "import time:"
    for line in result.stderr.decode().splitlines():
        if line.startswith(LINE_PREFIX):
            line = line[len(LINE_PREFIX):]
            self_time, total_time, depth_name = line.split(" | ")
            try:
                self_time = int(self_time) // 1000
                total_time = int(total_time) // 1000
            except:
                # importtime header will fail to parse
                continue
            # figure out depth from spaces preceding module name
            name = depth_name.lstrip()
            depth = (len(depth_name) - len(name) - 1) // 2 + 1

            start = start_stack[-1]
            # probably going up one level to parent
            while len(start_stack) > depth:
                start = start_stack.pop()
            # not sure how far down we're going...
            while len(start_stack) < depth:
                start_stack.append(start)
            # save new start time
            start_stack.append(start + total_time)

            # never skip the specified module(s)
            if name not in args.module:
                if args.min_elapsed and total_time < args.min_elapsed:
                    continue
                if args.max_depth is not None and depth > args.max_depth:
                    continue

            imports.append([name, start, self_time, total_time, depth])

    # write results
    title = f"importtime: {args.module} (total={int(runtime)}ms)"
    imports_to_table(imports, title, args.module, runtime, args.width)


if __name__ == "__main__":
    main()