1#! /usr/bin/env python3
2
3"""Tool for measuring execution time of small code snippets.
4
5This module avoids a number of common traps for measuring execution
6times.  See also Tim Peters' introduction to the Algorithms chapter in
7the Python Cookbook, published by O'Reilly.
8
9Library usage: see the Timer class.
10
11Command line usage:
12    python timeit.py [-n N] [-r N] [-s S] [-p] [-h] [--] [statement]
13
14Options:
15  -n/--number N: how many times to execute 'statement' (default: see below)
16  -r/--repeat N: how many times to repeat the timer (default 5)
17  -s/--setup S: statement to be executed once initially (default 'pass').
18                Execution time of this setup statement is NOT timed.
19  -p/--process: use time.process_time() (default is time.perf_counter())
20  -v/--verbose: print raw timing results; repeat for more digits precision
21  -u/--unit: set the output time unit (nsec, usec, msec, or sec)
22  -h/--help: print this usage message and exit
23  --: separate options from statement, use when statement starts with -
24  statement: statement to be timed (default 'pass')
25
26A multi-line statement may be given by specifying each line as a
27separate argument; indented lines are possible by enclosing an
28argument in quotes and using leading spaces.  Multiple -s options are
29treated similarly.
30
31If -n is not given, a suitable number of loops is calculated by trying
32increasing numbers from the sequence 1, 2, 5, 10, 20, 50, ... until the
33total time is at least 0.2 seconds.
34
35Note: there is a certain baseline overhead associated with executing a
36pass statement.  It differs between versions.  The code here doesn't try
37to hide it, but you should be aware of it.  The baseline overhead can be
38measured by invoking the program without arguments.
39
40Classes:
41
42    Timer
43
44Functions:
45
46    timeit(string, string) -> float
47    repeat(string, string) -> list
48    default_timer() -> float
49
50"""
51
52import gc
53import sys
54import time
55import itertools
56
57__all__ = ["Timer", "timeit", "repeat", "default_timer"]
58
59dummy_src_name = "<timeit-src>"
60default_number = 1000000
61default_repeat = 5
62default_timer = time.perf_counter
63
64_globals = globals
65
66# Don't change the indentation of the template; the reindent() calls
67# in Timer.__init__() depend on setup being indented 4 spaces and stmt
68# being indented 8 spaces.
69template = """
70def inner(_it, _timer{init}):
71    {setup}
72    _t0 = _timer()
73    for _i in _it:
74        {stmt}
75    _t1 = _timer()
76    return _t1 - _t0
77"""
78
79def reindent(src, indent):
80    """Helper to reindent a multi-line statement."""
81    return src.replace("\n", "\n" + " "*indent)
82
83class Timer:
84    """Class for timing execution speed of small code snippets.
85
86    The constructor takes a statement to be timed, an additional
87    statement used for setup, and a timer function.  Both statements
88    default to 'pass'; the timer function is platform-dependent (see
89    module doc string).  If 'globals' is specified, the code will be
90    executed within that namespace (as opposed to inside timeit's
91    namespace).
92
93    To measure the execution time of the first statement, use the
94    timeit() method.  The repeat() method is a convenience to call
95    timeit() multiple times and return a list of results.
96
97    The statements may contain newlines, as long as they don't contain
98    multi-line string literals.
99    """
100
101    def __init__(self, stmt="pass", setup="pass", timer=default_timer,
102                 globals=None):
103        """Constructor.  See class doc string."""
104        self.timer = timer
105        local_ns = {}
106        global_ns = _globals() if globals is None else globals
107        init = ''
108        if isinstance(setup, str):
109            # Check that the code can be compiled outside a function
110            compile(setup, dummy_src_name, "exec")
111            stmtprefix = setup + '\n'
112            setup = reindent(setup, 4)
113        elif callable(setup):
114            local_ns['_setup'] = setup
115            init += ', _setup=_setup'
116            stmtprefix = ''
117            setup = '_setup()'
118        else:
119            raise ValueError("setup is neither a string nor callable")
120        if isinstance(stmt, str):
121            # Check that the code can be compiled outside a function
122            compile(stmtprefix + stmt, dummy_src_name, "exec")
123            stmt = reindent(stmt, 8)
124        elif callable(stmt):
125            local_ns['_stmt'] = stmt
126            init += ', _stmt=_stmt'
127            stmt = '_stmt()'
128        else:
129            raise ValueError("stmt is neither a string nor callable")
130        src = template.format(stmt=stmt, setup=setup, init=init)
131        self.src = src  # Save for traceback display
132        code = compile(src, dummy_src_name, "exec")
133        exec(code, global_ns, local_ns)
134        self.inner = local_ns["inner"]
135
136    def print_exc(self, file=None):
137        """Helper to print a traceback from the timed code.
138
139        Typical use:
140
141            t = Timer(...)       # outside the try/except
142            try:
143                t.timeit(...)    # or t.repeat(...)
144            except:
145                t.print_exc()
146
147        The advantage over the standard traceback is that source lines
148        in the compiled template will be displayed.
149
150        The optional file argument directs where the traceback is
151        sent; it defaults to sys.stderr.
152        """
153        import linecache, traceback
154        if self.src is not None:
155            linecache.cache[dummy_src_name] = (len(self.src),
156                                               None,
157                                               self.src.split("\n"),
158                                               dummy_src_name)
159        # else the source is already stored somewhere else
160
161        traceback.print_exc(file=file)
162
163    def timeit(self, number=default_number):
164        """Time 'number' executions of the main statement.
165
166        To be precise, this executes the setup statement once, and
167        then returns the time it takes to execute the main statement
168        a number of times, as a float measured in seconds.  The
169        argument is the number of times through the loop, defaulting
170        to one million.  The main statement, the setup statement and
171        the timer function to be used are passed to the constructor.
172        """
173        it = itertools.repeat(None, number)
174        gcold = gc.isenabled()
175        gc.disable()
176        try:
177            timing = self.inner(it, self.timer)
178        finally:
179            if gcold:
180                gc.enable()
181        return timing
182
183    def repeat(self, repeat=default_repeat, number=default_number):
184        """Call timeit() a few times.
185
186        This is a convenience function that calls the timeit()
187        repeatedly, returning a list of results.  The first argument
188        specifies how many times to call timeit(), defaulting to 5;
189        the second argument specifies the timer argument, defaulting
190        to one million.
191
192        Note: it's tempting to calculate mean and standard deviation
193        from the result vector and report these.  However, this is not
194        very useful.  In a typical case, the lowest value gives a
195        lower bound for how fast your machine can run the given code
196        snippet; higher values in the result vector are typically not
197        caused by variability in Python's speed, but by other
198        processes interfering with your timing accuracy.  So the min()
199        of the result is probably the only number you should be
200        interested in.  After that, you should look at the entire
201        vector and apply common sense rather than statistics.
202        """
203        r = []
204        for i in range(repeat):
205            t = self.timeit(number)
206            r.append(t)
207        return r
208
209    def autorange(self, callback=None):
210        """Return the number of loops and time taken so that total time >= 0.2.
211
212        Calls the timeit method with increasing numbers from the sequence
213        1, 2, 5, 10, 20, 50, ... until the time taken is at least 0.2
214        second.  Returns (number, time_taken).
215
216        If *callback* is given and is not None, it will be called after
217        each trial with two arguments: ``callback(number, time_taken)``.
218        """
219        i = 1
220        while True:
221            for j in 1, 2, 5:
222                number = i * j
223                time_taken = self.timeit(number)
224                if callback:
225                    callback(number, time_taken)
226                if time_taken >= 0.2:
227                    return (number, time_taken)
228            i *= 10
229
230def timeit(stmt="pass", setup="pass", timer=default_timer,
231           number=default_number, globals=None):
232    """Convenience function to create Timer object and call timeit method."""
233    return Timer(stmt, setup, timer, globals).timeit(number)
234
235def repeat(stmt="pass", setup="pass", timer=default_timer,
236           repeat=default_repeat, number=default_number, globals=None):
237    """Convenience function to create Timer object and call repeat method."""
238    return Timer(stmt, setup, timer, globals).repeat(repeat, number)
239
240def main(args=None, *, _wrap_timer=None):
241    """Main program, used when run as a script.
242
243    The optional 'args' argument specifies the command line to be parsed,
244    defaulting to sys.argv[1:].
245
246    The return value is an exit code to be passed to sys.exit(); it
247    may be None to indicate success.
248
249    When an exception happens during timing, a traceback is printed to
250    stderr and the return value is 1.  Exceptions at other times
251    (including the template compilation) are not caught.
252
253    '_wrap_timer' is an internal interface used for unit testing.  If it
254    is not None, it must be a callable that accepts a timer function
255    and returns another timer function (used for unit testing).
256    """
257    if args is None:
258        args = sys.argv[1:]
259    import getopt
260    try:
261        opts, args = getopt.getopt(args, "n:u:s:r:tcpvh",
262                                   ["number=", "setup=", "repeat=",
263                                    "time", "clock", "process",
264                                    "verbose", "unit=", "help"])
265    except getopt.error as err:
266        print(err)
267        print("use -h/--help for command line help")
268        return 2
269
270    timer = default_timer
271    stmt = "\n".join(args) or "pass"
272    number = 0 # auto-determine
273    setup = []
274    repeat = default_repeat
275    verbose = 0
276    time_unit = None
277    units = {"nsec": 1e-9, "usec": 1e-6, "msec": 1e-3, "sec": 1.0}
278    precision = 3
279    for o, a in opts:
280        if o in ("-n", "--number"):
281            number = int(a)
282        if o in ("-s", "--setup"):
283            setup.append(a)
284        if o in ("-u", "--unit"):
285            if a in units:
286                time_unit = a
287            else:
288                print("Unrecognized unit. Please select nsec, usec, msec, or sec.",
289                    file=sys.stderr)
290                return 2
291        if o in ("-r", "--repeat"):
292            repeat = int(a)
293            if repeat <= 0:
294                repeat = 1
295        if o in ("-p", "--process"):
296            timer = time.process_time
297        if o in ("-v", "--verbose"):
298            if verbose:
299                precision += 1
300            verbose += 1
301        if o in ("-h", "--help"):
302            print(__doc__, end=' ')
303            return 0
304    setup = "\n".join(setup) or "pass"
305
306    # Include the current directory, so that local imports work (sys.path
307    # contains the directory of this script, rather than the current
308    # directory)
309    import os
310    sys.path.insert(0, os.curdir)
311    if _wrap_timer is not None:
312        timer = _wrap_timer(timer)
313
314    t = Timer(stmt, setup, timer)
315    if number == 0:
316        # determine number so that 0.2 <= total time < 2.0
317        callback = None
318        if verbose:
319            def callback(number, time_taken):
320                msg = "{num} loop{s} -> {secs:.{prec}g} secs"
321                plural = (number != 1)
322                print(msg.format(num=number, s='s' if plural else '',
323                                  secs=time_taken, prec=precision))
324        try:
325            number, _ = t.autorange(callback)
326        except:
327            t.print_exc()
328            return 1
329
330        if verbose:
331            print()
332
333    try:
334        raw_timings = t.repeat(repeat, number)
335    except:
336        t.print_exc()
337        return 1
338
339    def format_time(dt):
340        unit = time_unit
341
342        if unit is not None:
343            scale = units[unit]
344        else:
345            scales = [(scale, unit) for unit, scale in units.items()]
346            scales.sort(reverse=True)
347            for scale, unit in scales:
348                if dt >= scale:
349                    break
350
351        return "%.*g %s" % (precision, dt / scale, unit)
352
353    if verbose:
354        print("raw times: %s" % ", ".join(map(format_time, raw_timings)))
355        print()
356    timings = [dt / number for dt in raw_timings]
357
358    best = min(timings)
359    print("%d loop%s, best of %d: %s per loop"
360          % (number, 's' if number != 1 else '',
361             repeat, format_time(best)))
362
363    best = min(timings)
364    worst = max(timings)
365    if worst >= best * 4:
366        import warnings
367        warnings.warn_explicit("The test results are likely unreliable. "
368                               "The worst time (%s) was more than four times "
369                               "slower than the best time (%s)."
370                               % (format_time(worst), format_time(best)),
371                               UserWarning, '', 0)
372    return None
373
374if __name__ == "__main__":
375    sys.exit(main())
376