batchgrid 0.2.0b2__tar.gz → 0.2.0b3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: batchgrid
-Version: 0.2.0b2
+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
@@ -112,9 +112,10 @@ sheet as the tab last showed it. For a DataFrame, that is a copy of yours with t
 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:
 
@@ -194,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.0b2 → batchgrid-0.2.0b3}/README.md

@@ -86,9 +86,10 @@ sheet as the tab last showed it. For a DataFrame, that is a copy of yours with t
 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:
 
@@ -168,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.0b2 → batchgrid-0.2.0b3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "batchgrid"
-version = "0.2.0b2"
+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.0b2 → 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.0b2"
+__version__ = "0.2.0b3"
 
 __all__ = [
     "BatchgridError",

{batchgrid-0.2.0b2 → 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.0b2 → batchgrid-0.2.0b3}/src/batchgrid/_ui.py

@@ -3,9 +3,11 @@
 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
@@ -40,10 +42,12 @@ def ui(
     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.
     """
     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):
@@ -51,8 +55,15 @@ def ui(
         try:
             sheet = os.path.join(out.get(), "sheet.csv")
             args: List[str] = ["ui", *([path] if path else []), "--output", sheet, "--json"]
-            if embed or 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)]
 
@@ -65,6 +76,8 @@ 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":
@@ -120,6 +133,31 @@ 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:

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

@@ -169,6 +169,44 @@ def test_in_a_notebook_the_workspace_shows_in_the_cell(monkeypatch: pytest.Monke
     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