implement some timing wrappers

hbw/util.py

@@ -6,8 +6,10 @@
 
 from __future__ import annotations
 
-from typing import Hashable, Iterable, Callable
 
+import time
+from typing import Hashable, Iterable, Callable
+from functools import wraps
 import tracemalloc
 
 import law
@@ -211,6 +213,7 @@ def call_once_on_config(include_hash=False):
     Parametrized decorator to ensure that function *func* is only called once for the config *config*
     """
     def outer(func):
+        @wraps(func)
         def inner(config, *args, **kwargs):
             tag = f"{func.__name__}_called"
             if include_hash:
@@ -223,3 +226,31 @@ def inner(config, *args, **kwargs):
             return func(config, *args, **kwargs)
         return inner
     return outer
+
+
+def timeit(func):
+    """ Simple wrapper to measure execution time of a function """
+    @wraps(func)
+    def timeit_wrapper(*args, **kwargs):
+        start_time = time.perf_counter()
+        result = func(*args, **kwargs)
+        end_time = time.perf_counter()
+        total_time = end_time - start_time
+        _logger.info(f"Function {func.__name__}{args} {kwargs} Took {total_time:.4f} seconds")
+        return result
+    return timeit_wrapper
+
+
+def timeit_multiple(func):
+    """ Wrapper to measure the number of execution calls and the added execution time of a function """
+    @wraps(func)
+    def timeit_wrapper(*args, **kwargs):
+        func.total_calls = getattr(func, "total_calls", 0) + 1
+        start_time = time.perf_counter()
+        result = func(*args, **kwargs)
+        end_time = time.perf_counter()
+        total_time = end_time - start_time
+        func.total_time = getattr(func, "total_time", 0) + total_time
+        _logger.info(f"{func.__name__} has been run {func.total_calls} times ({func.total_time:.4f} seconds)")
+        return result
+    return timeit_wrapper