Add InfernoProfiler

README.md

@@ -11,7 +11,7 @@ pip install pyinferno
 ```
 ## Usage
 ### Context manager
-To profile a specific piece of code, use the [pyinstrument Profiler class](https://pyinstrument.readthedocs.io/en/latest/guide.html#profile-a-specific-chunk-of-code), then render the result using pyinferno's `InfernoRenderer`:
+To profile a specific piece of code, you can use the [pyinstrument Profiler class](https://pyinstrument.readthedocs.io/en/latest/guide.html#profile-a-specific-chunk-of-code), then render the result using pyinferno's `InfernoRenderer`:
 ```python
 from pyinstrument.profiler import Profiler
 from pyinferno import InfernoRenderer
@@ -23,12 +23,24 @@ def slow():
 with Profiler() as profiler:
     slow()
 
-output = profiler.render(InfernoRenderer(title="slow"))
+output = profiler.output(InfernoRenderer(title="slow"))
 
 with open("flamegraph.svg", "w+") as f:
     f.write(output)
 ```
 
+For convenience, the same result can be achieved using the `InfernoProfiler` context manager:
+```python
+from pyinferno import InfernoProfiler
+import time
+
+def slow():
+    time.sleep(0.5)
+
+with InfernoProfiler(file="flamegraph.svg", auto_open=True, title="slow"):
+    slow()
+```
+
 ### Command-line
 To profile a Python script, you can pass `pyinferno.Renderer` as the renderer to the `pyinstrument` CLI:
 ```bash

pyinferno/__init__.py

@@ -1,5 +1,6 @@
 from .pyinferno import InfernoError, flamegraph_from_lines
 from .renderer import InfernoRenderer as Renderer
+from .profiler import InfernoProfiler
 
 InfernoRenderer = Renderer
 

pyinferno/profiler.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import contextlib
+import os
+import tempfile
+import webbrowser
+from pathlib import Path
+
+from pyinstrument.profiler import Profiler
+
+from .renderer import InfernoRenderer
+
+
+@contextlib.contextmanager
+def InfernoProfiler(
+    file: str | Path | None = None,
+    auto_open: bool = False,
+    title: str | None = None,
+    **kwargs,
+):
+    if file is None and not auto_open:
+        raise ValueError(
+            "Must pass a file to save the flamegraph to, or auto_open=True to open directly in the browser."
+        )
+    profiler = Profiler(**kwargs)
+    profiler.start()
+    try:
+        yield profiler
+    finally:
+        profiler.stop()
+        output = profiler.output(InfernoRenderer(title=title))
+
+        if file:
+            outpath = file
+        else:
+            outpath = tempfile.NamedTemporaryFile(prefix="flamegraph-", suffix=".svg", delete=False).name
+
+        with open(outpath, "w+") as f:
+            f.write(output)
+
+        if auto_open:
+            filepath = os.path.relpath(outpath)
+            webbrowser.open(filepath)

pyproject.toml

@@ -15,7 +15,6 @@ requires-python = ">=3.7"
 classifiers = [
     "Programming Language :: Rust",
     "Programming Language :: Python :: Implementation :: CPython",
-    "Framework :: Pytest",
 ]
 dependencies = [
     "pyinstrument"

tests/test_pyinferno.py

@@ -6,7 +6,7 @@
 import pytest
 from pyinstrument.profiler import Profiler
 
-from pyinferno import InfernoError, Renderer, flamegraph_from_lines
+from pyinferno import InfernoError, Renderer, flamegraph_from_lines, InfernoProfiler
 
 TITLE = "My fantastic title"
 
@@ -49,14 +49,30 @@ def test_pyinstrument_default_title():
     assert len(result) > 0
 
 
-def test_profiler_context_manager():
+def test_pyinstrument_context_manager():
     with Profiler() as p:
         time.sleep(0.2)
     result = p.output(Renderer(title=TITLE))
     assert len(result) > 0
     assert TITLE in result
 
 
+def test_pyinferno_context_manager(tmp_path):
+    out_path = tmp_path / "out.svg"
+    with InfernoProfiler(file=out_path, title=TITLE) as p:
+        time.sleep(0.2)
+    with open(out_path) as f:
+        result = f.read()
+    assert len(result) > 0
+    assert TITLE in result
+
+
+def test_pyinferno_context_manager_raises_with_no_args():
+    with pytest.raises(ValueError):
+        with InfernoProfiler() as p:
+            time.sleep(0.2)
+
+
 def test_pyinstrument_cli(tmp_path):
     code_path = tmp_path / "code.py"
     out_path = tmp_path / "out.svg"