Add small utility to profile any function

profiling/profiling.py

@@ -0,0 +1,29 @@
+import functools
+
+from pyinstrument.profiler import Profiler
+
+
+def profile_function(output_file="profile.html"):
+    """
+    Profiles a function execution time.
+
+    Parameters
+    ----------
+    output_file: file to write profile output. Defaults to "profile.html".
+    """
+
+    def decorator(function):
+        @functools.wraps(function)
+        def wrapper(*args, **kwargs):
+            profiler = Profiler()
+            profiler.start()
+            result = function(*args, **kwargs)
+            profiler.stop()
+            output = profiler.output_html()
+            with open(output_file, "w") as f:
+                f.write(output)
+            return result
+
+        return wrapper
+
+    return decorator

test_requirements.txt

@@ -7,3 +7,4 @@ coverage>=6.4.4
 flake8>=3.9.2
 isort>=5.8.0
 mypy>=0.740
+pyinstrument>=4.4.0