batchgrid 0.2.0b1__tar.gz → 0.2.0b3__tar.gz

{batchgrid-0.2.0b1 → batchgrid-0.2.0b3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: batchgrid
-Version: 0.2.0b1
+Version: 0.2.0b3
 Summary: Run a prompt, or a whole pipeline, over every row of a DataFrame or spreadsheet - with cost estimates, retries and resumable runs.
 Project-URL: Homepage, https://github.com/mertguvencli/batchgrid
 Project-URL: Source, https://github.com/mertguvencli/batchgrid/tree/main/python
@@ -83,7 +83,8 @@ The file is the same one `batchgrid --save-plan` writes and `batchgrid --plan` r
 
 ## In the browser
 
-`batchgrid.ui()` opens the batchgrid workspace in your browser with your data in the sheet. You can
+`batchgrid.ui()` opens the batchgrid workspace with your data in the sheet. In a Jupyter notebook it
+shows in the cell's output; elsewhere, or with `inline=False`, it opens in your browser. You can
 chat, check the plan and run it there, while the log prints in the notebook cell or terminal as it
 happens:
 
@@ -105,14 +106,16 @@ batchgrid: 100%|██████████| 1000/1000 [00:42<00:00, failed=0
 Finished in the browser.
 ```
 
-To come back, press **Finish** in the tab, close the tab, or interrupt the cell. `ui()` returns the
+To come back, press **Finish** in the workspace, close its tab, or interrupt the cell. Once it stops,
+a note takes the workspace's place in the cell. `ui()` returns the
 sheet as the tab last showed it. For a DataFrame, that is a copy of yours with the new columns and
 any cells you edited; untouched columns keep their dtypes. `ui()` with no data starts on an empty
 sheet.
 
-The workspace runs on your machine and uses the keys from `batchgrid config` or your environment. Your
-browser has to be able to reach the machine running Python, so this does not work for a notebook on
-a remote server.
+The workspace runs on the machine running Python and uses the keys from `batchgrid config` or the
+environment there. Your browser has to be able to reach that machine. On Google Colab it reaches it
+through Colab's own proxy, so `ui()` works there too; a notebook on another remote server is out of
+reach.
 
 From a terminal, the `batchgrid` command this package installs does the same:
 
@@ -163,9 +166,10 @@ except KeyboardInterrupt:
     result = batchgrid.resume(df)
 ```
 
-### `batchgrid.ui(data=None, *, open_browser=True, port=None, log=True) -> DataFrame | None`
+### `batchgrid.ui(data=None, *, inline=None, height=720, open_browser=True, port=None, log=True) -> DataFrame | None`
 
-See [In the browser](#in-the-browser). `on_event` receives the same events as in `run()`, plus `ui`,
+See [In the browser](#in-the-browser). `inline` shows the workspace in the cell, `height` pixels tall;
+by default it does in a Jupyter notebook and opens a browser tab anywhere else. `on_event` receives the same events as in `run()`, plus `ui`,
 `sheet`, `browser_error` and `closed`.
 
 ### Errors
@@ -191,7 +195,7 @@ for the CLI in this order:
 1. the `cli=[...]` argument
 2. the `BATCHGRID_CLI` environment variable, e.g. `node /path/to/cli/dist/index.js`
 3. the npm `batchgrid` on the PATH (never this package's own command)
-4. `npx "batchgrid@>=0.3.0 <1"`
+4. `npx "batchgrid@>=0.3.1 <1"`
 
 ## Development
 

{batchgrid-0.2.0b1 → batchgrid-0.2.0b3}/README.md

@@ -57,7 +57,8 @@ The file is the same one `batchgrid --save-plan` writes and `batchgrid --plan` r
 
 ## In the browser
 
-`batchgrid.ui()` opens the batchgrid workspace in your browser with your data in the sheet. You can
+`batchgrid.ui()` opens the batchgrid workspace with your data in the sheet. In a Jupyter notebook it
+shows in the cell's output; elsewhere, or with `inline=False`, it opens in your browser. You can
 chat, check the plan and run it there, while the log prints in the notebook cell or terminal as it
 happens:
 
@@ -79,14 +80,16 @@ batchgrid: 100%|██████████| 1000/1000 [00:42<00:00, failed=0
 Finished in the browser.
 ```
 
-To come back, press **Finish** in the tab, close the tab, or interrupt the cell. `ui()` returns the
+To come back, press **Finish** in the workspace, close its tab, or interrupt the cell. Once it stops,
+a note takes the workspace's place in the cell. `ui()` returns the
 sheet as the tab last showed it. For a DataFrame, that is a copy of yours with the new columns and
 any cells you edited; untouched columns keep their dtypes. `ui()` with no data starts on an empty
 sheet.
 
-The workspace runs on your machine and uses the keys from `batchgrid config` or your environment. Your
-browser has to be able to reach the machine running Python, so this does not work for a notebook on
-a remote server.
+The workspace runs on the machine running Python and uses the keys from `batchgrid config` or the
+environment there. Your browser has to be able to reach that machine. On Google Colab it reaches it
+through Colab's own proxy, so `ui()` works there too; a notebook on another remote server is out of
+reach.
 
 From a terminal, the `batchgrid` command this package installs does the same:
 
@@ -137,9 +140,10 @@ except KeyboardInterrupt:
     result = batchgrid.resume(df)
 ```
 
-### `batchgrid.ui(data=None, *, open_browser=True, port=None, log=True) -> DataFrame | None`
+### `batchgrid.ui(data=None, *, inline=None, height=720, open_browser=True, port=None, log=True) -> DataFrame | None`
 
-See [In the browser](#in-the-browser). `on_event` receives the same events as in `run()`, plus `ui`,
+See [In the browser](#in-the-browser). `inline` shows the workspace in the cell, `height` pixels tall;
+by default it does in a Jupyter notebook and opens a browser tab anywhere else. `on_event` receives the same events as in `run()`, plus `ui`,
 `sheet`, `browser_error` and `closed`.
 
 ### Errors
@@ -165,7 +169,7 @@ for the CLI in this order:
 1. the `cli=[...]` argument
 2. the `BATCHGRID_CLI` environment variable, e.g. `node /path/to/cli/dist/index.js`
 3. the npm `batchgrid` on the PATH (never this package's own command)
-4. `npx "batchgrid@>=0.3.0 <1"`
+4. `npx "batchgrid@>=0.3.1 <1"`
 
 ## Development
 

{batchgrid-0.2.0b1 → batchgrid-0.2.0b3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "batchgrid"
-version = "0.2.0b1"
+version = "0.2.0b3"
 description = "Run a prompt, or a whole pipeline, over every row of a DataFrame or spreadsheet - with cost estimates, retries and resumable runs."
 readme = "README.md"
 license = "MIT"

{batchgrid-0.2.0b1 → batchgrid-0.2.0b3}/src/batchgrid/__init__.py

@@ -16,7 +16,7 @@ from ._ui import ui
 from ._cli import CLI_SPEC, find_cli
 from .errors import BatchgridError, CostLimitError, MissingSecretsError, NoPlanError
 
-__version__ = "0.2.0b1"
+__version__ = "0.2.0b3"
 
 __all__ = [
     "BatchgridError",

{batchgrid-0.2.0b1 → batchgrid-0.2.0b3}/src/batchgrid/_cli.py

@@ -19,8 +19,8 @@ import threading
 from collections import deque
 from typing import Callable, Deque, Dict, Iterator, List, Mapping, Optional, Sequence
 
-#: The CLI releases this package speaks to: 0.3.0 was the first with ``ui``, and its events only gain fields
-CLI_SPEC = "batchgrid@>=0.3.0 <1"
+#: The CLI releases this package speaks to: 0.3.0 was the first with ``ui``, 0.3.1 the first behind Colab's proxy, and its events only gain fields
+CLI_SPEC = "batchgrid@>=0.3.1 <1"
 
 Event = Dict[str, object]
 

{batchgrid-0.2.0b1 → batchgrid-0.2.0b3}/src/batchgrid/_ui.py

@@ -1,11 +1,13 @@
-"""ui() - the web workspace in the browser, with its log here and its sheet coming back."""
+"""ui() - the web workspace in the notebook or the browser, with its log here and its sheet coming back."""
 
 from __future__ import annotations
 
 import os
+import socket
 import sys
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, List, Mapping, Optional, Sequence
+from urllib.parse import urlsplit, urlunsplit
 
 from ._api import Data, EventHandler, _input_file, _is_dataframe, _merge_output, _TempDir
 from ._cli import Event, find_cli, stream_events
@@ -18,6 +20,8 @@ if TYPE_CHECKING:
 def ui(
     data: Optional[Data] = None,
     *,
+    inline: Optional[bool] = None,
+    height: int = 720,
     open_browser: bool = True,
     port: Optional[int] = None,
     log: bool = True,
@@ -25,28 +29,41 @@ def ui(
     cli: Optional[Sequence[str]] = None,
     env: Optional[Mapping[str, str]] = None,
 ) -> Optional["pd.DataFrame"]:
-    """Open the batchgrid workspace in the browser and wait until you are done there.
+    """Open the batchgrid workspace and wait until you are done there.
 
-    ``data`` - a file path or a pandas DataFrame - opens in the sheet. Chat,
-    plan and run in the browser; what happens there is printed here as it
-    happens, with a progress bar while rows run.
+    ``data`` - a file path or a pandas DataFrame - opens in the sheet. In a
+    Jupyter notebook the workspace shows in the cell's output, ``height``
+    pixels tall; ``inline=False`` opens it in a browser tab instead, which is
+    what happens outside a notebook. Chat, plan and run there; what happens is
+    printed here as it happens, with a progress bar while rows run.
 
-    Press Finish in the tab, close it, or interrupt the cell (Ctrl+C) to come
-    back. The sheet as the tab last showed it is returned as a DataFrame: for
+    Press Finish in the workspace, close its tab, or interrupt the cell
+    (Ctrl+C) to come back. The sheet as the tab last showed it is returned as a DataFrame: for
     a DataFrame you passed in, a copy with the new columns added and your
     edits in it. None when nothing was opened, or pandas is not installed.
 
-    It runs on this machine, so it needs a browser that can reach it - not a
-    notebook running on a remote server.
+    It runs where Python runs, so the browser has to reach that machine. On
+    Colab it does through Colab's own proxy; a notebook on another remote
+    server is out of reach.
     """
-    printer = _Printer() if log else None
+    embed = _in_notebook() if inline is None else inline
+    colab = _in_colab()
+    frame_view = _Frame(height) if embed else None
+    printer = _Printer(link=not embed) if log else None
     with _input_file(data) as (path, frame):
         out = _TempDir()
         try:
             sheet = os.path.join(out.get(), "sheet.csv")
             args: List[str] = ["ui", *([path] if path else []), "--output", sheet, "--json"]
-            if not open_browser:
+            # On Colab the browser that would open is on Google's machine, not yours
+            if embed or colab or not open_browser:
                 args.append("--no-open")
+            proxy: Optional[str] = None
+            if colab:
+                # The proxy's address names the port, so the port is settled before the server starts
+                port = port if port is not None else _free_port()
+                proxy = _colab_proxy(port)
+                args += ["--allow-host", str(urlsplit(proxy).hostname)]
             if port is not None:
                 args += ["--port", str(port)]
 
@@ -59,6 +76,10 @@ def ui(
                         failure = event
                     elif event.get("type") == "ui":
                         started = True
+                        if proxy is not None:
+                            event = {**event, "url": _through(proxy, str(event.get("url")))}
+                        if frame_view is not None:
+                            frame_view.show(str(event.get("url")))
                     elif event.get("type") == "exit":
                         ended = event
                     if printer is not None:
@@ -71,6 +92,8 @@ def ui(
             finally:
                 if printer is not None:
                     printer.close()
+                if frame_view is not None:
+                    frame_view.close()
             if failure is not None:
                 raise BatchgridError(str(failure.get("code")), str(failure.get("message")))
             if not started and ended is not None and ended.get("code"):
@@ -110,11 +133,81 @@ def _read_sheet(frame: Any, sheet: Path, sent: Optional[Path]) -> Optional["pd.D
     return merged
 
 
+def _in_colab() -> bool:
+    """Whether this is a Colab kernel, on a machine the browser reaches only through Colab's proxy"""
+    return "google.colab" in sys.modules
+
+
+def _colab_proxy(port: int) -> str:
+    """The address Colab's proxy gives a port on the kernel's machine"""
+    from google.colab.output import eval_js  # type: ignore[import-not-found]
+
+    return str(eval_js(f"google.colab.kernel.proxyPort({port}, {{cache: false}})"))
+
+
+def _through(proxy: str, url: str) -> str:
+    """The server's link, with the proxy's address in place of 127.0.0.1"""
+    base, own = urlsplit(proxy), urlsplit(url)
+    query = "&".join(q for q in (base.query, own.query) if q)
+    return urlunsplit((base.scheme, base.netloc, own.path or "/", query, ""))
+
+
+def _free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
+        probe.bind(("127.0.0.1", 0))
+        return int(probe.getsockname()[1])
+
+
+def _in_notebook() -> bool:
+    """Whether this runs in a Jupyter kernel, whose front end can show the workspace in a cell"""
+    try:
+        from IPython import get_ipython
+
+        shell = get_ipython()
+        return shell is not None and "IPKernelApp" in shell.config
+    except Exception:  # noqa: BLE001 - no IPython means a plain console
+        return False
+
+
+class _Frame:
+    """The workspace in the cell's output; once it stops, a note takes its place"""
+
+    def __init__(self, height: int) -> None:
+        self.height = height
+        self.handle: Any = None
+
+    def show(self, url: str) -> None:
+        from IPython.display import IFrame, display
+
+        self.handle = display(
+            IFrame(
+                url,
+                width="100%",
+                height=self.height,
+                extras=[
+                    'style="border: 1px solid rgba(128, 128, 128, 0.35); border-radius: 6px"',
+                    'allow="clipboard-read; clipboard-write"',
+                ],
+            ),
+            display_id=True,
+        )
+
+    def close(self) -> None:
+        if self.handle is None:
+            return
+        from IPython.display import HTML
+
+        # The server behind the frame is gone, so a saved notebook keeps no dead workspace
+        self.handle.update(HTML("<em>batchgrid closed. Run the cell again to reopen it.</em>"))
+        self.handle = None
+
+
 class _Printer:
     """The terminal's log, in a notebook cell or a Python console"""
 
-    def __init__(self) -> None:
+    def __init__(self, link: bool = True) -> None:
         self.bar: Any = None
+        self.link = link
         try:
             from tqdm.auto import tqdm
 
@@ -130,7 +223,7 @@ class _Printer:
         # Anything else ends the run's bar, so its lines are not drawn over it
         if kind in ("result", "closed", "user", "plan"):
             self.close()
-        if kind == "ui":
+        if kind == "ui" and self.link:
             self._link(str(event.get("url")))
         elif kind == "user":
             self._line(f"\n❯ {event.get('text')}")
@@ -155,7 +248,7 @@ class _Printer:
         elif kind == "error":
             self._line(f"✖ {event.get('message')}")
         elif kind == "closed":
-            why = {"tab_closed": "The tab was closed", "finished": "Finished in the browser"}.get(
+            why = {"tab_closed": "The tab was closed", "finished": "Finished in the browser" if self.link else "Finished"}.get(
                 str(event.get("reason")), "Stopped"
             )
             self._line(f"\n{why}.")
@@ -173,15 +266,11 @@ class _Printer:
         self.bar.refresh()
 
     def _link(self, url: str) -> None:
-        try:
-            from IPython import get_ipython
+        if _in_notebook():
             from IPython.display import HTML, display
 
-            if get_ipython() is not None and "IPKernelApp" in get_ipython().config:
-                display(HTML(f'batchgrid is open at <a href="{url}" target="_blank">{url.split("?")[0]}</a>'))
-                return
-        except Exception:  # noqa: BLE001 - no IPython means a plain console
-            pass
+            display(HTML(f'batchgrid is open at <a href="{url}" target="_blank">{url.split("?")[0]}</a>'))
+            return
         self._line(f"batchgrid is open at {url}")
 
     def _line(self, text: str) -> None:

{batchgrid-0.2.0b1 → batchgrid-0.2.0b3}/tests/test_ui.py

@@ -14,6 +14,7 @@ import stat
 import sys
 import threading
 import time
+import types
 import urllib.request
 from pathlib import Path
 from typing import Any, Dict, List
@@ -129,6 +130,83 @@ def test_interrupting_the_cell_returns_instead_of_raising() -> None:
     assert result.loc[0, "sentiment"] == "positive"
 
 
+def test_in_a_notebook_the_workspace_shows_in_the_cell(monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]) -> None:
+    shown: List[Any] = []
+
+    class Handle:
+        def update(self, obj: Any) -> None:
+            shown.append(obj)
+
+    def display(obj: Any, display_id: bool = False) -> Handle:
+        shown.append(obj)
+        return Handle()
+
+    # A stand-in for the IPython a Jupyter kernel runs
+    ipython = types.ModuleType("IPython")
+    ipython.get_ipython = lambda: types.SimpleNamespace(config={"IPKernelApp": {}})  # type: ignore[attr-defined]
+    ipython_display = types.ModuleType("IPython.display")
+    ipython_display.HTML = lambda html: html  # type: ignore[attr-defined]
+    ipython_display.IFrame = lambda src, **options: {"src": src, **options}  # type: ignore[attr-defined]
+    ipython_display.display = display  # type: ignore[attr-defined]
+    monkeypatch.setitem(sys.modules, "IPython", ipython)
+    monkeypatch.setitem(sys.modules, "IPython.display", ipython_display)
+
+    urls: List[str] = []
+
+    def act(tab: FakeTab) -> None:
+        urls.append(tab.url)
+        tab.call("finish", {})
+
+    # No open_browser=False: showing it in the cell never opens a tab
+    batchgrid.ui(on_event=open_tab(act))
+
+    assert len(shown) == 2
+    assert shown[0]["src"] == urls[0] and shown[0]["height"] == 720
+    # The server behind it stopped, so a note takes the workspace's place
+    assert "batchgrid closed" in shown[1]
+    printed = capsys.readouterr().out
+    assert "is open at" not in printed
+    assert "Finished." in printed
+
+
+def test_on_colab_the_workspace_goes_through_its_proxy(monkeypatch: pytest.MonkeyPatch) -> None:
+    asked: List[str] = []
+
+    def eval_js(script: str) -> str:
+        asked.append(script)
+        port = script.split("(")[1].split(",")[0]
+        return f"https://{port}-abc-colab.googleusercontent.com/?authuser=0"
+
+    # A stand-in for the module a Colab kernel has loaded
+    colab = types.ModuleType("google.colab")
+    colab_output = types.ModuleType("google.colab.output")
+    colab_output.eval_js = eval_js  # type: ignore[attr-defined]
+    monkeypatch.setitem(sys.modules, "google.colab", colab)
+    monkeypatch.setitem(sys.modules, "google.colab.output", colab_output)
+
+    seen: Dict[str, Any] = {}
+
+    def on_event(event: Dict[str, Any]) -> None:
+        if event["type"] != "ui":
+            return
+        seen["url"] = event["url"]
+        port = asked[0].split("(")[1].split(",")[0]
+        token = event["url"].split("token=")[1]
+        # What Colab's proxy sends on: the request addressed to its own host name
+        page = urllib.request.Request(
+            f"http://127.0.0.1:{port}/?token={token}", headers={"Host": f"{port}-abc-colab.googleusercontent.com"}
+        )
+        with urllib.request.urlopen(page, timeout=10) as response:
+            seen["status"] = response.status
+        tab = FakeTab(f"http://127.0.0.1:{port}/?token={token}")
+        threading.Thread(target=tab.call, args=("finish", {}), daemon=True).start()
+
+    batchgrid.ui(inline=False, log=False, on_event=on_event)
+
+    assert seen["url"].startswith("https://") and "-abc-colab.googleusercontent.com/?authuser=0&token=" in seen["url"]
+    assert seen["status"] == 200
+
+
 def test_nothing_synced_returns_none() -> None:
     result = batchgrid.ui(open_browser=False, log=False, on_event=open_tab(lambda tab: tab.call("finish", {})))
     assert result is None