X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;f=stopwatch.py;h=ae8e01aa18452ed932ef589b9047c7f0c68ab58a;hb=b8d28ce7d3d913975131284aeaa5d5caf951a631;hp=1326cb1fec8ffb41db461f99aac59372ded655f1;hpb=4faa994d32223c8d560d9dad0ca90a3f7eb10d6a;p=python_utils.git

diff --git a/stopwatch.py b/stopwatch.py
index 1326cb1..ae8e01a 100644
--- a/stopwatch.py
+++ b/stopwatch.py
@@ -1,27 +1,37 @@
 #!/usr/bin/env python3
 
+"""A simple stopwatch decorator / context for timing things."""
+
+import contextlib
 import time
-from typing import Callable
+from typing import Callable, Literal
 
 
-class Timer(object):
+class Timer(contextlib.AbstractContextManager):
     """
-    with timer.Timer() as t:
-        do_the_thing()
+    A stopwatch to time how long something takes (walltime).
+
+    e.g.
 
-    walltime = t()
-    print(f'That took {walltime}s.')
+        with stopwatch.Timer() as t:
+            do_the_thing()
+
+        walltime = t()
+        print(f'That took {walltime} seconds.')
     """
 
     def __init__(self) -> None:
-        self.start = None
-        self.end = 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) -> bool:
+    def __exit__(self, *args) -> Literal[False]:
         self.end = time.perf_counter()
-        return None  # don't suppress exceptions
+        return False