From c33086ebf6207822c9faa48889737fb315849777 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Janek=20Nouvertn=C3=A9?=
 <25355197+provinzkraut@users.noreply.github.com>
Date: Wed, 4 Oct 2023 19:17:24 +0200
Subject: [PATCH] add enable_profiling contextmanager

---
 README.md                   |  7 +++++++
 scalene/scalene_profiler.py | 10 +++++++++-
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 67ff64974..f339d729d 100755
--- a/README.md
+++ b/README.md
@@ -108,6 +108,13 @@ scalene_profiler.start()
 scalene_profiler.stop()
 ```
 
+```Python
+from scalene.scalene_profiler import enable_profiling
+
+with enable_profiling():
+    # do something
+```
+
 </details>
 
 <details>
diff --git a/scalene/scalene_profiler.py b/scalene/scalene_profiler.py
index 230b6c994..92fe3cda2 100644
--- a/scalene/scalene_profiler.py
+++ b/scalene/scalene_profiler.py
@@ -62,7 +62,7 @@ def nada(*args):
     Tuple,
     Type,
     Union,
-    cast,
+    cast, Generator,
 )
 
 from scalene.scalene_arguments import ScaleneArguments
@@ -160,6 +160,14 @@ def stop() -> None:
     Scalene.stop()
 
 
+@contextlib.contextmanager
+def enable_profiling() -> Generator[None, None, None]:
+    """Contextmanager that starts and stops profiling"""
+    start()
+    yield
+    stop()
+
+
 def _get_module_details(
     mod_name: str,
     error: Type[Exception] = ImportError,