feat: marimo edit --sandbox nb.py (#2134)

* feat: marimo edit --sandbox nb.py

* add tests

* fix

* fix again

* fix

* Update marimo/_cli/cli.py

* Update marimo/_cli/cli.py

* fix

---------

Co-authored-by: Akshay Agrawal <[email protected]>

marimo/_cli/cli.py

@@ -4,6 +4,7 @@
 import json
 import os
 import pathlib
+import sys
 import tempfile
 from typing import Any, Optional, get_args
 
@@ -255,6 +256,17 @@ def main(log_level: str, quiet: bool, development_mode: bool) -> None:
     type=bool,
     help="Don't check if a new version of marimo is available for download.",
 )
+@click.option(
+    "--sandbox",
+    is_flag=True,
+    default=False,
+    show_default=True,
+    type=bool,
+    help="""
+    Run the command in an isolated virtual environment using
+    'uv run --isolated'. Requires 'uv'.
+    """,
+)
 @click.option("--profile-dir", default=None, type=str, hidden=True)
 @click.argument("name", required=False)
 @click.argument("args", nargs=-1, type=click.UNPROCESSED)
@@ -268,10 +280,17 @@ def edit(
     base_url: str,
     allow_origins: Optional[tuple[str, ...]],
     skip_update_check: bool,
+    sandbox: bool,
     profile_dir: Optional[str],
     name: Optional[str],
     args: tuple[str, ...],
 ) -> None:
+    if sandbox:
+        from marimo._cli.sandbox import run_in_sandbox
+
+        run_in_sandbox(sys.argv[1:], name)
+        return
+
     GLOBAL_SETTINGS.PROFILE_DIR = profile_dir
     if not skip_update_check and os.getenv("MARIMO_SKIP_UPDATE_CHECK") != "1":
         GLOBAL_SETTINGS.CHECK_STATUS_UPDATE = True
@@ -501,6 +520,17 @@ def new(
     type=bool,
     help="Redirect console logs to the browser console.",
 )
+@click.option(
+    "--sandbox",
+    is_flag=True,
+    default=False,
+    show_default=True,
+    type=bool,
+    help="""
+    Run the command in an isolated virtual environment using
+    'uv run --isolated'. Requires `uv`.
+    """,
+)
 @click.argument("name", required=True)
 @click.argument("args", nargs=-1, type=click.UNPROCESSED)
 def run(
@@ -515,9 +545,16 @@ def run(
     base_url: str,
     allow_origins: tuple[str, ...],
     redirect_console_to_browser: bool,
+    sandbox: bool,
     name: str,
     args: tuple[str, ...],
 ) -> None:
+    if sandbox:
+        from marimo._cli.sandbox import run_in_sandbox
+
+        run_in_sandbox(sys.argv[1:], name)
+        return
+
     # Validate name, or download from URL
     # The second return value is an optional temporary directory. It is unused,
     # but must be kept around because its lifetime on disk is bound to the life

marimo/_cli/sandbox.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+import os
+import re
+import subprocess
+from typing import Any, Dict, List, Optional, cast
+
+from marimo import _loggers
+from marimo._dependencies.dependencies import DependencyManager
+
+LOGGER = _loggers.marimo_logger()
+
+REGEX = (
+    r"(?m)^# /// (?P<type>[a-zA-Z0-9-]+)$\s(?P<content>(^#(| .*)$\s)+)^# ///$"
+)
+
+
+def run_in_sandbox(
+    args: List[str],
+    name: Optional[str] = None,
+) -> subprocess.CompletedProcess[Any]:
+    import click
+
+    if not DependencyManager.which("uv"):
+        raise click.UsageError("uv must be installed to use --sandbox")
+
+    cmd = ["marimo"] + args
+    cmd.remove("--sandbox")
+
+    # If name if a filepath, parse the dependencies from the file
+    dependencies = []
+    if name is not None and os.path.isfile(name):
+        with open(name) as f:
+            dependencies = _get_dependencies(f.read()) or []
+        # Add marimo, if it's not already there
+        if "marimo" not in dependencies and len(dependencies) > 0:
+            dependencies.append("marimo")
+
+    if dependencies:
+        import tempfile
+
+        with tempfile.NamedTemporaryFile(
+            mode="w", delete=False, suffix=".txt"
+        ) as temp_file:
+            temp_file.write("\n".join(dependencies))
+            temp_file_path = temp_file.name
+
+        cmd = [
+            "uv",
+            "run",
+            "--isolated",
+            "--with-requirements",
+            temp_file_path,
+        ] + cmd
+
+        # Clean up the temporary file after the subprocess has run
+        import atexit
+
+        atexit.register(lambda: os.unlink(temp_file_path))
+    else:
+        cmd = ["uv", "run", "--isolated"] + cmd
+
+    click.echo(f"Running in a sandbox: {' '.join(cmd)}")
+
+    return subprocess.run(cmd)
+
+
+def _get_dependencies(script: str) -> List[str] | None:
+    try:
+        pyproject = _read_pyproject(script) or {}
+        return cast(List[str], pyproject.get("dependencies", []))
+    except Exception as e:
+        LOGGER.warning(f"Failed to parse dependencies: {e}")
+        return None
+
+
+def _read_pyproject(script: str) -> Dict[str, Any] | None:
+    """
+    Read the pyproject.toml file from the script.
+
+    Adapted from https://peps.python.org/pep-0723/#reference-implementation
+    """
+    name = "script"
+    matches = list(
+        filter(lambda m: m.group("type") == name, re.finditer(REGEX, script))
+    )
+    if len(matches) > 1:
+        raise ValueError(f"Multiple {name} blocks found")
+    elif len(matches) == 1:
+        content = "".join(
+            line[2:] if line.startswith("# ") else line[1:]
+            for line in matches[0].group("content").splitlines(keepends=True)
+        )
+        import tomlkit
+
+        return tomlkit.parse(content)
+    else:
+        return None

tests/_cli/test_cli.py

@@ -24,8 +24,11 @@
 from marimo import __version__
 from marimo._ast import codegen
 from marimo._ast.cell import CellConfig
+from marimo._dependencies.dependencies import DependencyManager
 from marimo._utils.config.config import ROOT_DIR as CONFIG_ROOT_DIR
 
+HAS_UV = DependencyManager.which("uv")
+
 
 def _is_win32() -> bool:
     return sys.platform == "win32"
@@ -493,6 +496,43 @@ def test_cli_custom_host() -> None:
     _check_contents(p, b"marimo-mode data-mode='edit'", contents)
 
 
+@pytest.mark.skipif(not HAS_UV, reason="uv is required for sandbox tests")
+def test_cli_sandbox_edit(temp_marimo_file: str) -> None:
+    port = _get_port()
+    p = subprocess.Popen(
+        [
+            "marimo",
+            "edit",
+            temp_marimo_file,
+            "-p",
+            str(port),
+            "--headless",
+            "--no-token",
+            "--sandbox",
+        ]
+    )
+    contents = _try_fetch(port)
+    _check_contents(p, b"marimo-mode data-mode='edit'", contents)
+
+
+@pytest.mark.skipif(not HAS_UV, reason="uv is required for sandbox tests")
+def test_cli_sandbox_run(temp_marimo_file: str) -> None:
+    port = _get_port()
+    p = subprocess.Popen(
+        [
+            "marimo",
+            "run",
+            temp_marimo_file,
+            "-p",
+            str(port),
+            "--headless",
+            "--sandbox",
+        ]
+    )
+    contents = _try_fetch(port)
+    _check_contents(p, b"marimo-mode data-mode='read'", contents)
+
+
 @pytest.mark.xfail(condition=_is_win32(), reason="flaky on Windows")
 def test_cli_edit_interrupt_twice() -> None:
     # two SIGINTs should kill the CLI

tests/_cli/test_sandbox.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from marimo._cli.sandbox import _get_dependencies
+
+
+def test_get_dependencies():
+    SCRIPT = """
+# Copyright 2024 Marimo. All rights reserved.
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#     "polars",
+#     "marimo>=0.8.0",
+#     "quak",
+#     "vega-datasets",
+# ]
+# ///
+
+import marimo
+
+__generated_with = "0.8.2"
+app = marimo.App(width="medium")
+"""
+    assert _get_dependencies(SCRIPT) == [
+        "polars",
+        "marimo>=0.8.0",
+        "quak",
+        "vega-datasets",
+    ]
+
+
+def test_no_dependencies():
+    SCRIPT = """
+import marimo
+
+__generated_with = "0.8.2"
+app = marimo.App(width="medium")
+"""
+    assert _get_dependencies(SCRIPT) == []