add enable_profiling contextmanager

provinzkraut · Oct 4, 2023 · c33086e · c33086e
1 parent 20ed6b6
commit c33086e
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git 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
@@ -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,