src/pyutils/stopwatch.py

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