Easy profiling option

setup.cfg

@@ -114,6 +114,8 @@ doc =
     pydantic==1.10
     quartodoc==0.7.2
     griffe==0.33.0
+profile =
+    pyinstrument
 
 [options.packages.find]
 include = shiny, shiny.*

shiny/_main.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import contextlib
 import copy
 import importlib
 import importlib.util
@@ -20,6 +21,7 @@
 
 from . import _autoreload, _hostenv, _static, _utils
 from ._docstring import no_example
+from ._profiler import check_profiler_dependencies, profiler
 from ._typing_extensions import NotRequired, TypedDict
 from .express import is_express_app
 from .express._utils import escape_to_var_name
@@ -150,6 +152,15 @@ def main() -> None:
     help="Dev mode",
     show_default=True,
 )
+@click.option(
+    "--profile",
+    is_flag=True,
+    default=False,
+    help="Run the app in profiling mode. This will launch the app under a "
+    "sampling profiler, and after Shiny is stopped, open a Speedscope.app flamegraph "
+    "in a web browser.",
+    show_default=True,
+)
 @no_example()
 def run(
     app: str | shiny.App,
@@ -167,8 +178,17 @@ def run(
     factory: bool,
     launch_browser: bool,
     dev_mode: bool,
+    profile: bool,
     **kwargs: object,
 ) -> None:
+    if profile:
+        if reload:
+            print(
+                "Error: --profile and --reload cannot be used together", file=sys.stderr
+            )
+            sys.exit(1)
+        check_profiler_dependencies()
+
     reload_includes_list = reload_includes.split(",")
     reload_excludes_list = reload_excludes.split(",")
     return run_app(
@@ -186,6 +206,7 @@ def run(
         factory=factory,
         launch_browser=launch_browser,
         dev_mode=dev_mode,
+        profile=profile,
         **kwargs,
     )
 
@@ -206,6 +227,7 @@ def run_app(
     factory: bool = False,
     launch_browser: bool = False,
     dev_mode: bool = True,
+    profile: bool = False,
     **kwargs: object,
 ) -> None:
     """
@@ -250,6 +272,10 @@ def run_app(
         Treat ``app`` as an application factory, i.e. a () -> <ASGI app> callable.
     launch_browser
         Launch app browser after app starts, using the Python webbrowser module.
+    profile
+        Run the app in profiling mode. This will launch the app under a sampling
+        profiler, and after Shiny is stopped, open a Speedscope.app flamegraph in a web
+        browser.
     **kwargs
         Additional keyword arguments which are passed to ``uvicorn.run``. For more
         information see [Uvicorn documentation](https://www.uvicorn.org/).
@@ -353,19 +379,20 @@ def run_app(
 
     maybe_setup_rsw_proxying(log_config)
 
-    uvicorn.run(  # pyright: ignore[reportUnknownMemberType]
-        app,
-        host=host,
-        port=port,
-        ws_max_size=ws_max_size,
-        log_level=log_level,
-        log_config=log_config,
-        app_dir=app_dir,
-        factory=factory,
-        lifespan="on",
-        **reload_args,  # pyright: ignore[reportArgumentType]
-        **kwargs,
-    )
+    with profiler() if profile else contextlib.nullcontext():
+        uvicorn.run(  # pyright: ignore[reportUnknownMemberType]
+            app,
+            host=host,
+            port=port,
+            ws_max_size=ws_max_size,
+            log_level=log_level,
+            log_config=log_config,
+            app_dir=app_dir,
+            factory=factory,
+            lifespan="on",
+            **reload_args,  # pyright: ignore[reportArgumentType]
+            **kwargs,
+        )
 
 
 def setup_hot_reload(

shiny/_profiler.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+import contextlib
+import sys
+from typing import Any, Generator
+
+
+def check_profiler_dependencies() -> None:
+    try:
+        import pyinstrument  # pyright: ignore[reportUnusedImport]
+    except ImportError:
+        print(
+            "Error: Profiler is not installed. You can install it with "
+            "'pip install shiny[profile]'.",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+
+@contextlib.contextmanager
+def profiler() -> Generator[None, Any, None]:
+    import base64
+    import os
+    import time
+    import webbrowser
+    from urllib.parse import quote_plus
+
+    import pyinstrument
+
+    prof = pyinstrument.Profiler()
+    prof.start()
+
+    epoch_time = int(time.time())
+    output_filename = f"profile-{epoch_time}.json"
+    output_filename_abs = os.path.join(os.getcwd(), output_filename)
+    print(f"Profiling to {output_filename_abs}", file=sys.stderr)
+
+    try:
+        yield
+    finally:
+        prof_session = prof.stop()
+
+        import pyinstrument.renderers.speedscope
+
+        renderer = pyinstrument.renderers.speedscope.SpeedscopeRenderer()
+        output_str = renderer.render(prof_session)
+
+        with open(output_filename_abs, "w", encoding="utf-8") as f:
+            f.write(output_str)
+
+        print(f"Profile saved to {output_filename_abs}", file=sys.stderr)
+
+        b64_str = base64.b64encode(output_str.encode("utf-8")).decode("utf-8")
+        data_uri = f"data:application/json;base64,{b64_str}"
+        full_url = (
+            "https://speedscope.app/#profileURL="
+            + quote_plus(data_uri)
+            # + "&title="
+            # + quote_plus(output_filename)
+        )
+
+        webbrowser.open(full_url)