stopwatch.py

   1 #!/usr/bin/env python3
   2
   3 """A simple stopwatch decorator / context for timing things."""
   4
   5 import contextlib
   6 import time
   7 from typing import Callable, Literal
   8
   9
  10 class Timer(contextlib.AbstractContextManager):
  11     """
  12     A stopwatch to time how long something takes (walltime).
  13
  14     e.g.
  15
  16         with stopwatch.Timer() as t:
  17             do_the_thing()
  18
  19         walltime = t()
  20         print(f'That took {walltime} seconds.')
  21     """
  22
  23     def __init__(self) -> None:
  24         self.start = 0.0
  25         self.end = 0.0
  26
  27     def __enter__(self) -> Callable[[], float]:
  28         """Returns a functor that, when called, returns the walltime of the
  29         operation in seconds.
  30         """
  31         self.start = time.perf_counter()
  32         self.end = 0.0
  33         return lambda: self.end - self.start
  34
  35     def __exit__(self, *args) -> Literal[False]:
  36         self.end = time.perf_counter()
  37         return False