Timeit asv benchmark setup() routines, look for slow ones.

timeit_asv_setup.py

"""
Timeit asv benchmark setup() routines, look for slow ones.

Looks for benchmarks under './benchmarks/' directory.
"""
import os
import time
import itertools
import tempfile
import shutil
import re

from asv.benchmark import disc_benchmarks


def main():
    cwd = os.getcwd()

    benchmarks = dict((b.name, b) for b in disc_benchmarks('benchmarks'))

    setups = {}
    for benchmark in benchmarks.values():
        for setup in benchmark._setups:
            m = re.match('<bound method (.*) of <(.*) object.*>>', repr(setup))
            if m:
                name = m.group(2) + ':' + m.group(1)
            else:
                name = repr(setup)
            setups.setdefault(name, []).append(benchmark)

    slow_setups = {}

    for setup, benchmarks in sorted(setups.items()):
        count = len(benchmarks)
        benchmark = benchmarks[0]

        tmpdir = tempfile.mkdtemp()
        try:
            os.chdir(tmpdir)

            for j, p in enumerate(itertools.product(*benchmark.params)):
                if benchmark.params:
                    name = setup + '(' + ','.join(p) + ')'
                    benchmark.set_param_idx(j)
                else:
                    name = setup

                if benchmark._setup_cache:
                    try:
                        start = time.perf_counter()
                        benchmark.do_setup_cache()
                        stop = time.perf_counter()
                        setup_cache_time = stop - start
                    except Exception as exc:
                        setup_cache_time = exc
                else:
                    setup_cache_time = None

                try:
                    start = time.perf_counter()
                    benchmark.do_setup()
                    stop = time.perf_counter()

                    setup_time = stop - start
                except Exception as exc:
                    setup_time = exc

                repeats = 0
                for b in benchmarks:
                    repeat = getattr(b, 'repeat', 1)
                    try:
                        repeat = repeat[1]
                    except TypeError:
                        repeat = int(repeat)
                        if repeat == 0:
                            repeat = 10
                    repeats += repeat

                print("")
                print(name)
                print("    benchmarks:", len(benchmarks))
                print("    repeats:", repeats)
                print("    setup:", setup_time)

                
                if setup_cache_time is not None:
                    print("    cache:", setup_cache_time)

                if isinstance(setup_time, float):
                    total_time = max(count, repeats) * setup_time
                    if total_time > 5:
                        print("    SLOW!  (total {:.1f}s)".format(total_time))
                        slow_setups[name] = (total_time, setup_time, max(count, repeats))
        except KeyboardInterrupt:
            break
        finally:
            os.chdir(cwd)
            shutil.rmtree(tmpdir)

    print("")
    print("Slow setup routines")

    items = list(slow_setups.items())
    items.sort(key=lambda x: (x[1][0], x[0]))

    for name, (total_time, setup_time, count) in items:
        print("{}: {:.2g} * {:.2g} = {:.1f}s".format(name, setup_time, count, total_time))


if __name__ == "__main__":
    main()