criterion-polyglot 0.1.0

#!/usr/bin/env python3

import argparse as harness_argparse
import sys as harness_sys
import timeit as harness_timeit
from typing import Optional

@"imports"@

@"declarations"@

def main(timed: str, setup: Optional[str], init: Optional[str]):
    setup = 'pass' if setup is None else setup
    init = 'pass' if init is None else init

    # The init code can do whatever it wants, and any variables
    # it sets will be persisted in init_globals, which is passed
    # into timeit.Timer as the global scope
    init_globals = dict()
    exec(init, init_globals)

    bench = harness_timeit.Timer(
        stmt=timed,
        setup=setup,
        globals=init_globals
    )

    for line in map(str.rstrip, harness_sys.stdin):
        # TODO: check whether value is valid
        n = int(line)
        if n == 0:
            print("0 nsec", flush=True)
            harness_sys.exit(0)

        try:
            elapsed = bench.timeit(n)
        except Exception:
            bench.print_exc()
            harness_sys.exit(1)

        elapsed_nsec = round(elapsed * 1e9)
        print(f"{elapsed_nsec} nsec", flush=True)


def cli() -> harness_argparse.Namespace:
    parser = harness_argparse.ArgumentParser()
    parser.add_argument('filename', type=harness_argparse.FileType('r'))
    parser.add_argument('-s', '--setup', type=harness_argparse.FileType('r'))
    parser.add_argument('-S', '--init', type=harness_argparse.FileType('r'))
    return parser.parse_args()


if __name__ == "__main__":
    args = cli()
    timed = args.filename.read()
    init = args.init.read() if args.init is not None else None
    setup = args.setup.read() if args.setup is not None else None

    main(timed, setup, init)