Pretty print replacement for timeit

pp_timeit.py

# coding: utf-8
"""
Pretty print replacement for timeit

pp_timeit.auto_timeit runs a statement and automatically
determines the number of repetitions required to gain
satisfactory accuracy (similar to timeit.main).

pp_timeit uses auto_timeit and prints the result so that the
different magnitudes of results are visible by their alignment.

>>> from pp_timeit import pp_timeit, auto_timeit

>>> auto_timeit("1 + 1")
2.188011194448336e-08
>>> pp_timeit("1 + 1")
1 + 1                                           20 ns/pass (best of 90x524288)

>>> setup = '''
... import time
... def test(n):
...     time.sleep(10 ** -n)
... '''
>>> auto_timeit("test(5)", setup)
0.0026257429271936417

>>> _ = (
... pp_timeit("test(8)", setup),
... pp_timeit("test(6)", setup),
... pp_timeit("test(4)", setup),
... pp_timeit("test(2)", setup),
... pp_timeit("test(1)", setup),
... )
test(8)                                        658 ns/pass (best of 57x131072)
test(6)                                      76686 ns/pass (best of 62x1024)
test(4)                                     176    us/pass (best of 52x512)
test(2)                                   10323    us/pass (best of 59x8)
test(1)                                  100       ms/pass (best of 49x1)

The doctests are for illustration only, they obviously fail
because the run times are variable.
"""

from __future__ import print_function
from __future__ import unicode_literals

import sys
if sys.version_info.major < 3:
    range = xrange

from timeit import Timer


NS_PER_SEC = 1000000000
UNITS = [
    (1, " s"),
    (1000, "ms"),
    (1000000, "us"),
    (NS_PER_SEC, "ns"),
]

def _auto_timeit(stmt, setup="pass"):
    t = Timer(stmt, setup=setup)

    # determine number so that time per repeat > 50 ms
    for i in range(30):
        number = 2 ** i
        per_repeat = t.timeit(number)
        if per_repeat > 0.05:
            break

    # increase accuracy if we haven't already spent > 5 sec
    if per_repeat > 5:
        repeat = 1
        per_pass = per_repeat / number
    else:
        repeat = int(5 / per_repeat)
        res = t.repeat(repeat=repeat, number=number)
        per_pass = min(res) / number

    return per_pass, repeat, number


def auto_timeit(*args, **kwargs):
    per_pass, repeat, number = _auto_timeit(*args, **kwargs)
    return per_pass


def pp_timeit(stmt, *args, **kwargs):
    line_len = kwargs.pop('line_len', 80)
    name = kwargs.pop('name', None)
    if name is None:
        name = stmt
    max_name_len = line_len - 48
    name = name[:max_name_len]
    if len(name) == max_name_len:
        name += " "

    if 'test_args' in kwargs:
        per_pass, repeat, number = kwargs.pop('test_args')
    else:
        per_pass, repeat, number = _auto_timeit(stmt, *args, **kwargs)

    for i, unit in enumerate(UNITS):
        unit_val, unit_name = unit
        per_pass_in_unit = round(per_pass * unit_val)
        if per_pass_in_unit >= 100 or unit_name == 'ns':
            per_pass_str = "{0:.0f}".format(per_pass_in_unit)
            name_pad = line_len - (len(name) + 47)
            ns_num = (per_pass_in_unit / unit_val) * NS_PER_SEC
            ns_digits = str(int(ns_num))
            lead_pad = 17 - len(ns_digits)
            trail_pad = 3 * (3 - i) + 1
            pad = name_pad + lead_pad
            print(
                name + (" " * pad) +
                per_pass_str + (" " * trail_pad) +
                unit_name + "/pass (best of " +
                str(repeat) + "x" + str(number) + ")"
            )
            return


if __name__ == '__main__':
    from random import random
    rint = lambda : int(10 ** (random() * 5))

    pp_timeit("test", test_args=(100., rint(), rint()))
    pp_timeit("test", test_args=(99.999999, rint(), rint()))
    pp_timeit("test", test_args=(0.12, rint(), rint()))
    pp_timeit("test" * 100, test_args=(12345678.9, rint(), rint()))
    pp_timeit("test", test_args=(1234567.89, rint(), rint()))
    pp_timeit("test", test_args=(123456.789, rint(), rint()))
    pp_timeit("test", test_args=(12345.6789, rint(), rint()))
    pp_timeit("test", test_args=(1234.56789, rint(), rint()))
    pp_timeit("test", test_args=(123.456789, rint(), rint()))
    pp_timeit("test", test_args=(12.3456789, rint(), rint()))
    pp_timeit("test", test_args=(1.23456789, rint(), rint()))
    pp_timeit("test", test_args=(.123456789, rint(), rint()))
    pp_timeit("test", test_args=(.023456789, rint(), rint()))
    pp_timeit("test", test_args=(.003456789, rint(), rint()))
    pp_timeit("test", test_args=(.000456789, rint(), rint()))
    pp_timeit("test", test_args=(.000056789, rint(), rint()))
    pp_timeit("test", test_args=(.000006789, rint(), rint()))
    pp_timeit("test", test_args=(.000000089, rint(), rint()))
    pp_timeit("test", test_args=(.000000009, rint(), rint()))
    pp_timeit("test" * 100, test_args=(.000000001, rint(), rint()))