#!/usr/bin/env python3 """A simple stopwatch decorator / context for timing things.""" import time from typing import Callable, Optional class Timer(object): """ A stopwatch to time how long something takes (walltime). e.g. with stopwatch.Timer() as t: do_the_thing() walltime = t() print(f'That took {walltime} seconds.') """ def __init__(self) -> None: self.start = 0.0 self.end = 0.0 def __enter__(self) -> Callable[[], float]: """Returns a functor that, when called, returns the walltime of the operation in seconds. """ self.start = time.perf_counter() self.end = 0.0 return lambda: self.end - self.start def __exit__(self, *args) -> Optional[bool]: self.end = time.perf_counter() return None # don't suppress exceptions