From 8faeb8470f2511048678f927309a084b76278e6b Mon Sep 17 00:00:00 2001
From: Nodar Okroshiashvili <n.okroshiashvili@gmail.com>
Date: Thu, 2 Mar 2023 00:22:55 +0400
Subject: [PATCH] Add small utility to profile any function

---
 profiling/profiling.py | 29 +++++++++++++++++++++++++++++
 test_requirements.txt  |  1 +
 2 files changed, 30 insertions(+)
 create mode 100644 profiling/profiling.py

diff --git a/profiling/profiling.py b/profiling/profiling.py
new file mode 100644
index 000000000..78df3edd8
--- /dev/null
+++ b/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
diff --git a/test_requirements.txt b/test_requirements.txt
index 42ebfeb9b..d2c088546 100644
--- a/test_requirements.txt
+++ b/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