From c1386ba696bff17e55280ee16c61026348f08269 Mon Sep 17 00:00:00 2001
From: Peter Sharpe <pds@mit.edu>
Date: Thu, 21 Dec 2023 17:25:35 +0100
Subject: [PATCH] add a context manager for timing functions

---
 aerosandbox/tools/code_benchmarking.py | 33 ++++++++++++++++++++++++--
 1 file changed, 31 insertions(+), 2 deletions(-)

diff --git a/aerosandbox/tools/code_benchmarking.py b/aerosandbox/tools/code_benchmarking.py
index 7bd1b059..bb291370 100644
--- a/aerosandbox/tools/code_benchmarking.py
+++ b/aerosandbox/tools/code_benchmarking.py
@@ -3,11 +3,38 @@
 import aerosandbox.numpy as np
 
 
+class Timer(object):
+    """
+    A context manager for timing things. Use it like this:
+
+    with Timer("My timer"):  # You can optionally give it a name
+        # Do stuff
+
+    Results are printed to stdout. You can access the runtime (in seconds) directly with:
+
+    with Timer("My timer") as t:
+        # Do stuff
+    print(t.runtime)
+    """
+
+    def __init__(self, name=None):
+        self.name = name
+
+    def __enter__(self):
+        self.t_start = time.perf_counter_ns()
+
+    def __exit__(self, type, value, traceback):
+        self.runtime = (time.perf_counter_ns() - self.t_start) / 1e9
+        if self.name:
+            print(f'[{self.name}]')
+        print(f'Elapsed: {runtime} sec')
+
+
 def time_function(
         func: Callable,
         repeats: int = None,
         desired_runtime: float = None,
-        runtime_reduction = np.min,
+        runtime_reduction=np.min,
 ) -> Tuple[float, Any]:
     """
     Runs a given callable and tells you how long it took to run, in seconds. Also returns the result of the function
@@ -70,9 +97,11 @@ def time_function_once():
         result
     )
 
+
 if __name__ == '__main__':
 
     def f():
         time.sleep(0.1)
 
-    print(time_function(f, desired_runtime=1))
\ No newline at end of file
+
+    print(time_function(f, desired_runtime=1))