Context manager API v2

.github/workflows/test.yml

@@ -66,6 +66,7 @@ jobs:
       run: |
             docker run --rm -v ${{ github.workspace }}:/ws:rw --workdir=/ws \
             -e QEMU_EMULATED=1 \
+            -e CI=1 \
             ${{ env.img }} \
             bash -exc '${{ env.py }} -m venv .env && \
             source .env/bin/activate && \

docs/guide.md

@@ -70,8 +70,27 @@ spent.
 
 ## Profile a specific chunk of code
 
-Pyinstrument also has a Python API. Just surround your code with Pyinstrument,
-like this:
+Pyinstrument also has a Python API. You can use a with-block, like this:
+
+```python
+import pyinstrument
+
+with pyinstrument.profile():
+    # code you want to profile
+```
+
+Or you can decorate a function/method, like this:
+
+```python
+import pyinstrument
+
+@pyinstrument.profile()
+def my_function():
+    # code you want to profile
+
+```
+
+There's also a lower-level API called Profiler, that's more flexible:
 
 ```python
 from pyinstrument import Profiler
@@ -82,7 +101,6 @@ profiler.start()
 # code you want to profile
 
 profiler.stop()
-
 profiler.print()
 ```
 
@@ -91,6 +109,8 @@ If you get "No samples were recorded." because your code executed in under
 value smaller than the default 0.001 (1 millisecond) like this:
 
 ```python
+pyinstrument.profile(interval=0.0001)
+# or,
 profiler = Profiler(interval=0.0001)
 ...
 ```
@@ -104,6 +124,7 @@ save this HTML for later, use
 {meth}`profiler.output_html() <pyinstrument.Profiler.output_html>`.
 
 ## Profile code in Jupyter/IPython
+
 Via [IPython magics](https://ipython.readthedocs.io/en/stable/interactive/magics.html),
 you can profile a line or a cell in IPython or Jupyter.
 

docs/reference.md

@@ -17,6 +17,49 @@ print a profile report to the console.
 The Python API is also available, for calling pyinstrument directly from
 Python and writing integrations with with other tools.
 
+### The `profile` function
+
+For example:
+
+```python
+with pyinstrument.profile():
+    time.sleep(1)
+```
+
+This will print something like:
+
+```
+pyinstrument ........................................
+.
+.  Block at testfile.py:2
+.
+.  1.000 <module>  testfile.py:1
+.  └─ 1.000 sleep  <built-in>
+.
+.....................................................
+```
+
+You can also use it as a function/method decorator, like this:
+
+```python
+@pyinstrument.profile()
+def my_function():
+    time.sleep(1)
+```
+
+```{eval-rst}
+.. function:: pyinstrument.profile(*, interval=0.001, async_mode="disabled", \
+                                 use_timing_thread=None, renderer=None, \
+                                 target_description=None)
+
+    Creates a context-manager or function decorator object, which profiles the given code and prints the output to stdout.
+
+    The *interval*, *async_mode* and *use_timing_thread* parameters are passed through to the underlying :class:`pyinstrument.Profiler` object.
+
+    You can pass a renderer to customise the output. By default, it uses a :class:`ConsoleRenderer <pyinstrument.renderers.ConsoleRenderer>` with `short_mode` set.
+
+```
+
 ### The Profiler object
 
 ```{eval-rst}

examples/context_api.py

@@ -0,0 +1,43 @@
+import os
+import pprint
+import sys
+import time
+
+pprint.pprint(sys.path)
+import pyinstrument
+
+
+@pyinstrument.profile()
+def main():
+    py_file_count = 0
+    py_file_size = 0
+
+    print("Start.")
+    print("scanning home dir...")
+
+    with pyinstrument.profile():
+        for dir_path, dirnames, filenames in os.walk(os.path.expanduser("~/Music")):
+            for filename in filenames:
+                file_path = os.path.join(dir_path, filename)
+                _, ext = os.path.splitext(file_path)
+                if ext == ".py":
+                    py_file_count += 1
+                    try:
+                        py_file_size += os.stat(file_path).st_size
+                    except:
+                        pass
+
+    print("There are {} python files on your system.".format(py_file_count))
+    print("Total size: {} kB".format(py_file_size / 1024))
+
+
+class A:
+    @pyinstrument.profile()
+    def foo(self):
+        time.sleep(0.1)
+
+
+if __name__ == "__main__":
+    a = A()
+    a.foo()
+    main()

noxfile.py

@@ -26,3 +26,10 @@ def livedocs(session):
     session.env["UV_PRERELEASE"] = "allow"
     session.install("-e", ".[docs]")
     session.run("make", "-C", "docs", "livehtml")
+
+
+@nox.session(python=False)
+def htmldev(session):
+    with session.chdir("html_renderer"):
+        session.run("npm", "install")
+        session.run("npm", "run", "dev")

pyinstrument/__init__.py

@@ -1,8 +1,9 @@
 import warnings
 
+from pyinstrument.context_manager import profile
 from pyinstrument.profiler import Profiler
 
-__all__ = ["__version__", "Profiler", "load_ipython_extension"]
+__all__ = ["__version__", "Profiler", "load_ipython_extension", "profile"]
 __version__ = "4.6.2"
 
 # enable deprecation warnings

pyinstrument/__main__.py

@@ -375,7 +375,7 @@ def store_and_consume_remaining(
             use_timing_thread=options.use_timing_thread,
         )
 
-        profiler.start()
+        profiler.start(target_description=f'Program: {" ".join(argv)}')
 
         try:
             sys.argv[:] = argv

pyinstrument/context_manager.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+import functools
+import inspect
+import sys
+import typing
+
+from pyinstrument.profiler import AsyncMode, Profiler
+from pyinstrument.renderers.base import Renderer
+from pyinstrument.renderers.console import ConsoleRenderer
+from pyinstrument.typing import Unpack
+from pyinstrument.util import file_supports_color, file_supports_unicode
+
+CallableVar = typing.TypeVar("CallableVar", bound=typing.Callable)
+
+
+class ProfileContextOptions(typing.TypedDict, total=False):
+    interval: float
+    async_mode: AsyncMode
+    use_timing_thread: bool | None
+    renderer: Renderer | None
+    target_description: str | None
+
+
+class ProfileContext:
+    options: ProfileContextOptions
+
+    def __init__(
+        self,
+        **kwargs: Unpack[ProfileContextOptions],
+    ):
+        profiler_options = {
+            "interval": kwargs.get("interval", 0.001),
+            # note- different async mode from the default, because it's easy
+            # to run multiple profilers at once using the decorator/context
+            # manager
+            "async_mode": kwargs.get("async_mode", "disabled"),
+            "use_timing_thread": kwargs.get("use_timing_thread", None),
+        }
+        self.profiler = Profiler(**profiler_options)
+        self.options = kwargs
+
+    @typing.overload
+    def __call__(self, func: CallableVar, /) -> CallableVar: ...
+    @typing.overload
+    def __call__(self, /, **kwargs: Unpack[ProfileContextOptions]) -> "ProfileContext": ...
+    def __call__(
+        self, func: typing.Callable | None = None, /, **kwargs: Unpack[ProfileContextOptions]
+    ):
+        if func is not None:
+
+            @functools.wraps(func)
+            def wrapper(*args, **kwargs):
+                target_description = self.options.get("target_description")
+                if target_description is None:
+                    target_description = f"Function {func.__qualname__} at {func.__code__.co_filename}:{func.__code__.co_firstlineno}"
+
+                with self(target_description=target_description):
+                    return func(*args, **kwargs)
+
+            return typing.cast(typing.Callable, wrapper)
+        else:
+            return ProfileContext(**{**self.options, **kwargs})
+
+    def __enter__(self):
+        if self.profiler.is_running:
+            raise RuntimeError(
+                "This profiler is already running - did you forget the brackets on pyinstrument.profile() ?"
+            )
+
+        caller_frame = inspect.currentframe().f_back  # type: ignore
+        assert caller_frame is not None
+        target_description = self.options.get("target_description")
+        if target_description is None:
+            target_description = "Block at {}:{}".format(
+                caller_frame.f_code.co_filename, caller_frame.f_lineno
+            )
+
+        self.profiler.start(
+            caller_frame=caller_frame,
+            target_description=target_description,
+        )
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        session = self.profiler.stop()
+
+        renderer = self.options.get("renderer")
+        f = sys.stderr
+
+        if renderer is None:
+            renderer = ConsoleRenderer(
+                color=file_supports_color(f),
+                unicode=file_supports_unicode(f),
+                short_mode=True,
+            )
+
+        f.write(renderer.render(session))
+
+
+class _Profile:
+    @typing.overload
+    def __call__(self, func: CallableVar, /) -> CallableVar: ...
+    @typing.overload
+    def __call__(self, /, **kwargs: Unpack[ProfileContextOptions]) -> "ProfileContext": ...
+    def __call__(
+        self, func: typing.Callable | None = None, /, **kwargs: Unpack[ProfileContextOptions]
+    ):
+        if func is not None:
+            return ProfileContext(**kwargs)(func)
+        else:
+            return ProfileContext(**kwargs)
+
+
+profile = _Profile()

pyinstrument/frame.py

@@ -199,6 +199,15 @@ def is_application_code(self) -> bool:
             if file_path.startswith("<ipython-input-"):
                 # lines typed at a console or in a notebook are app code
                 return True
+            elif file_path == "<string>" or file_path == "<stdin>":
+                # eval/exec is app code if started by a parent frame that is
+                # app code
+                if self.parent:
+                    return self.parent.is_application_code
+                else:
+                    # if this is the root frame, it must have been started
+                    # with -c, so it's app code
+                    return True
             else:
                 # otherwise, this is probably some library-internal code gen
                 return False
@@ -335,7 +344,6 @@ def __repr__(self):
 
 
 class FrameGroup:
-    _libraries: list[str] | None
     _frames: list[Frame]
     _exit_frames: list[Frame] | None
 
@@ -344,25 +352,9 @@ def __init__(self, root: Frame):
         self.id = str(uuid.uuid4())
         self._frames = []
         self._exit_frames = None
-        self._libraries = None
 
         self.add_frame(root)
 
-    @property
-    def libraries(self) -> list[str]:
-        if self._libraries is None:
-            libraries: list[str] = []
-
-            for frame in self.frames:
-                if frame.file_path_short:
-                    library = frame.file_path_short.split(os.sep)[0]
-                    library, _ = os.path.splitext(library)
-                    if library and library not in libraries:
-                        libraries.append(library)
-            self._libraries = libraries
-
-        return self._libraries
-
     @property
     def frames(self) -> Sequence[Frame]:
         return tuple(self._frames)

pyinstrument/processors.py

@@ -281,9 +281,6 @@ def is_runpy_frame(frame: Frame):
 
     result = result.children[0]
 
-    if not is_runpy_frame(result):
-        return frame
-
     # at this point we know we've matched the first few frames of a command
     # line invocation. We'll trim some runpy frames and return.