| 1 | # Copyright 2026 Google LLC |
| 2 | # |
| 3 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | # you may not use this file except in compliance with the License. |
| 5 | # You may obtain a copy of the License at |
| 6 | # |
| 7 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | # |
| 9 | # Unless required by applicable law or agreed to in writing, software |
| 10 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | # See the License for the specific language governing permissions and |
| 13 | # limitations under the License. |
| 14 | |
| 15 | import datetime |
| 16 | import nbformat |
| 17 | import os |
| 18 | import re |
| 19 | import sys |
| 20 | import typer |
| 21 | import uuid |
| 22 | from nbformat.v4 import new_output |
| 23 | from rich.console import Console |
| 24 | from typing import List, Optional |
| 25 | from typing_extensions import Annotated |
| 26 | |
| 27 | from colab_cli.runtime import ColabRuntime |
| 28 | from colab_cli.utils import handle_image, is_terminal_error, render_display_data |
| 29 | from colab_cli.console import connect_console |
| 30 | |
| 31 | _console = Console() |
| 32 | |
| 33 | TITLE_REGEX = re.compile(r"^\s*#\s*@title\s+(.*)", re.MULTILINE) |
| 34 | ENV_KEY_REGEX = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") |
| 35 | |
| 36 | |
| 37 | def is_stdin_tty(): |
| 38 | return sys.stdin.isatty() |
| 39 | |
| 40 | |
| 41 | def _parse_env_vars(env: Optional[List[str]]) -> dict[str, str]: |
| 42 | """Parse repeatable --env KEY=VALUE entries into an ordered mapping.""" |
| 43 | env_vars = {} |
| 44 | for item in env or []: |
| 45 | if "=" not in item: |
| 46 | typer.echo( |
| 47 | f"[colab] Invalid --env value {item!r}. Expected KEY=VALUE.", |
| 48 | err=True, |
| 49 | ) |
| 50 | raise typer.Exit(2) |
| 51 | |
| 52 | key, value = item.split("=", 1) |
| 53 | if not ENV_KEY_REGEX.fullmatch(key): |
| 54 | typer.echo( |
| 55 | f"[colab] Invalid --env key {key!r}. Expected a valid " |
| 56 | "environment variable name.", |
| 57 | err=True, |
| 58 | ) |
| 59 | raise typer.Exit(2) |
| 60 | |
| 61 | env_vars[key] = value |
| 62 | return env_vars |
| 63 | |
| 64 | |
| 65 | def _build_env_prelude(env_vars: dict[str, str]) -> str: |
| 66 | """Build Python source that sets environment variables in the remote kernel.""" |
| 67 | if not env_vars: |
| 68 | return "" |
| 69 | |
| 70 | lines = ["import os"] |
| 71 | lines.extend(f"os.environ[{key!r}] = {value!r}" for key, value in env_vars.items()) |
| 72 | return "\n".join(lines) + "\n" |
| 73 | |
| 74 | |
| 75 | def save_output(outputs, cell): |
| 76 | if cell is None: |
| 77 | return |
| 78 | |
| 79 | if not hasattr(cell, "outputs"): |
| 80 | cell.outputs = [] |
| 81 | else: |
| 82 | cell.outputs.clear() |
| 83 | |
| 84 | for out in outputs: |
| 85 | if out.get("output_type") == "stream": |
| 86 | cell.outputs.append( |
| 87 | new_output( |
| 88 | output_type="stream", |
| 89 | name=out.get("name", "stdout"), |
| 90 | text=out.get("text", ""), |
| 91 | ) |
| 92 | ) |
| 93 | elif "data" in out: |
| 94 | output_type = out.get("output_type", "display_data") |
| 95 | cell.outputs.append( |
| 96 | new_output( |
| 97 | output_type=output_type, |
| 98 | data=out["data"], |
| 99 | metadata=out.get("metadata", {}), |
| 100 | ) |
| 101 | ) |
| 102 | elif out.get("output_type") == "error": |
| 103 | cell.outputs.append( |
| 104 | new_output( |
| 105 | output_type="error", |
| 106 | ename=out.get("ename", "Error"), |
| 107 | evalue=out.get("evalue", ""), |
| 108 | traceback=out.get("traceback", []), |
| 109 | ) |
| 110 | ) |
| 111 | |
| 112 | |
| 113 | def display_output(out, output_image=None): |
| 114 | if out.get("output_type") == "stream": |
| 115 | stream = sys.stderr if out.get("name") == "stderr" else sys.stdout |
| 116 | stream.write(out.get("text", "")) |
| 117 | stream.flush() |
| 118 | elif "data" in out: |
| 119 | data = out["data"] |
| 120 | text = render_display_data(data) |
| 121 | if text is not None: |
| 122 | _console.print(text) |
| 123 | if png := data.get("image/png"): |
| 124 | handle_image(png, "image/png", target_path=output_image) |
| 125 | elif jpeg := data.get("image/jpeg"): |
| 126 | handle_image(jpeg, "image/jpeg", target_path=output_image) |
| 127 | elif out.get("output_type") == "error": |
| 128 | tb = out.get("traceback", []) |
| 129 | if tb: |
| 130 | sys.stderr.write("".join(tb) + "\n") |
| 131 | else: |
| 132 | ename = out.get("ename", "Error") |
| 133 | evalue = out.get("evalue", "") |
| 134 | sys.stderr.write(f"{ename}: {evalue}\n") |
| 135 | else: |
| 136 | # Ignore silent outputs like metadata or clear_output for streaming |
| 137 | pass |
| 138 | |
| 139 | |
| 140 | def exec_command( |
| 141 | session: Annotated[ |
| 142 | Optional[str], typer.Option("-s", "--session", help="Session name") |
| 143 | ] = None, |
| 144 | file: Annotated[ |
| 145 | Optional[str], typer.Option("-f", "--file", help="File to execute") |
| 146 | ] = None, |
| 147 | output_image: Annotated[ |
| 148 | Optional[str], typer.Option("--output-image", help="Path to save plot") |
| 149 | ] = None, |
| 150 | timeout: Annotated[ |
| 151 | Optional[float], |
| 152 | typer.Option("--timeout", help="Timeout in seconds for code execution"), |
| 153 | ] = 30.0, |
| 154 | env: Annotated[ |
| 155 | Optional[List[str]], |
| 156 | typer.Option( |
| 157 | "--env", |
| 158 | help=( |
| 159 | "Set an environment variable in the remote kernel as KEY=VALUE. " |
| 160 | "Repeat for multiple variables." |
| 161 | ), |
| 162 | ), |
| 163 | ] = None, |
| 164 | ): |
| 165 | """Execute code in a session""" |
| 166 | from colab_cli.common import state |
| 167 | |
| 168 | env_vars = _parse_env_vars(env) |
| 169 | name = state.resolve_session(session) |
| 170 | s = state.store.get(name) |
| 171 | if not s: |
| 172 | typer.echo(f"[colab] Session '{name}' not found.") |
| 173 | raise typer.Exit(1) |
| 174 | |
| 175 | code_blocks = [] |
| 176 | if file: |
| 177 | if file.endswith(".ipynb"): |
| 178 | typer.echo(f"[colab] Parsing notebook '{file}'...") |
| 179 | with open(file, "r", encoding="utf-8") as f: |
| 180 | nb = nbformat.read(f, as_version=4) |
| 181 | for cell in nb.cells: |
| 182 | # nbformat v4.5+ requires 'id' at the top level |
| 183 | if not hasattr(cell, "id") or not cell.id: |
| 184 | cell.id = str(uuid.uuid4()) |
| 185 | |
| 186 | if cell.cell_type == "code": |
| 187 | code_blocks.append( |
| 188 | {"code": cell.source, "id": cell.id, "cell": cell} |
| 189 | ) |
| 190 | else: |
| 191 | with open(file, "r") as f: |
| 192 | code_blocks.append({"code": f.read(), "id": None}) |
| 193 | else: |
| 194 | if is_stdin_tty(): |
| 195 | typer.echo("[colab] Error: No input provided. Pipe code or provide a file.") |
| 196 | raise typer.Exit(1) |
| 197 | code_blocks.append({"code": sys.stdin.read(), "id": None}) |
| 198 | |
| 199 | if not any(b["code"].strip() for b in code_blocks): |
| 200 | raise typer.Exit(0) |
| 201 | |
| 202 | def on_started(kid): |
| 203 | s.kernel_id = kid |
| 204 | state.store.add(s) |
| 205 | |
| 206 | def on_sess_started(sid): |
| 207 | s.session_id = sid |
| 208 | state.store.add(s) |
| 209 | |
| 210 | runtime = ColabRuntime( |
| 211 | s.url, |
| 212 | s.token, |
| 213 | kernel_id=s.kernel_id, |
| 214 | session_id=s.session_id, |
| 215 | on_kernel_started=on_started, |
| 216 | on_session_started=on_sess_started, |
| 217 | ) |
| 218 | try: |
| 219 | # Ensure we are in /content which is the standard Colab working directory |
| 220 | runtime.execute_code( |
| 221 | "import os; os.makedirs('/content', exist_ok=True); os.chdir('/content')" |
| 222 | ) |
| 223 | except Exception as e: |
| 224 | if is_terminal_error(e): |
| 225 | typer.echo( |
| 226 | f"[colab] Session '{name}' appears to be lost (404/401). Cleaning up." |
| 227 | ) |
| 228 | state.prune_session(name) |
| 229 | raise typer.Exit(1) |
| 230 | raise e |
| 231 | |
| 232 | try: |
| 233 | is_nb = file and file.endswith(".ipynb") |
| 234 | s.running = f"exec({file or 'stdin'})" |
| 235 | state.store.add(s) |
| 236 | |
| 237 | for i, block in enumerate(code_blocks): |
| 238 | code = _build_env_prelude(env_vars) + block["code"] |
| 239 | identifier = None |
| 240 | if is_nb: |
| 241 | title_match = TITLE_REGEX.search(code) |
| 242 | if title_match: |
| 243 | identifier = title_match.group(1).strip() |
| 244 | elif block.get("id"): |
| 245 | identifier = block["id"] |
| 246 | else: |
| 247 | identifier = "" |
| 248 | |
| 249 | identifier_str = f" - {identifier}" if identifier else "" |
| 250 | typer.echo( |
| 251 | f"[colab] Executing cell {i + 1}/{len(code_blocks)}{identifier_str}..." |
| 252 | ) |
| 253 | |
| 254 | s.last_execution = ( |
| 255 | file or "stdin", |
| 256 | identifier, |
| 257 | datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), |
| 258 | ) |
| 259 | state.store.add(s) |
| 260 | |
| 261 | outputs = runtime.execute_code( |
| 262 | code, |
| 263 | output_hook=lambda o: display_output(o, output_image), |
| 264 | timeout=timeout, |
| 265 | ) |
| 266 | if "cell" in block: |
| 267 | save_output(outputs, block["cell"]) |
| 268 | state.history.log_event( |
| 269 | name, |
| 270 | "execution", |
| 271 | { |
| 272 | "code": code, |
| 273 | "outputs": outputs, |
| 274 | "cell_index": i if len(code_blocks) > 1 else None, |
| 275 | "cell_id": block.get("id"), |
| 276 | }, |
| 277 | ) |
| 278 | finally: |
| 279 | s.running = None |
| 280 | state.store.add(s) |
| 281 | runtime.stop() |
| 282 | if file and file.endswith(".ipynb"): |
| 283 | output_file = os.path.splitext(file)[0] + "_output.ipynb" |
| 284 | typer.echo(f"[colab] Saving notebook with outputs to '{output_file}'...") |
| 285 | with open(output_file, "w", encoding="utf-8") as f: |
| 286 | nbformat.write(nb, f) |
| 287 | |
| 288 | |
| 289 | def repl( |
| 290 | session: Annotated[ |
| 291 | Optional[str], typer.Option("-s", "--session", help="Session name") |
| 292 | ] = None, |
| 293 | output_image: Annotated[ |
| 294 | Optional[str], typer.Option("--output-image", help="Path to save plot") |
| 295 | ] = None, |
| 296 | ): |
| 297 | """Start an interactive REPL""" |
| 298 | from colab_cli.common import state |
| 299 | |
| 300 | name = state.resolve_session(session) |
| 301 | s = state.store.get(name) |
| 302 | if not s: |
| 303 | typer.echo(f"[colab] Session '{name}' not found.") |
| 304 | raise typer.Exit(1) |
| 305 | |
| 306 | def on_started(kid): |
| 307 | s.kernel_id = kid |
| 308 | state.store.add(s) |
| 309 | |
| 310 | def on_sess_started(sid): |
| 311 | s.session_id = sid |
| 312 | state.store.add(s) |
| 313 | |
| 314 | runtime = ColabRuntime( |
| 315 | s.url, |
| 316 | s.token, |
| 317 | kernel_id=s.kernel_id, |
| 318 | session_id=s.session_id, |
| 319 | on_kernel_started=on_started, |
| 320 | on_session_started=on_sess_started, |
| 321 | ) |
| 322 | try: |
| 323 | # Ensure we are in /content which is the standard Colab working directory |
| 324 | runtime.execute_code( |
| 325 | "import os; os.makedirs('/content', exist_ok=True); os.chdir('/content')" |
| 326 | ) |
| 327 | except Exception as e: |
| 328 | if is_terminal_error(e): |
| 329 | typer.echo( |
| 330 | f"[colab] Session '{name}' appears to be lost (404/401). Cleaning up." |
| 331 | ) |
| 332 | state.prune_session(name) |
| 333 | raise typer.Exit(1) |
| 334 | raise e |
| 335 | |
| 336 | if not is_stdin_tty(): |
| 337 | code = sys.stdin.read() |
| 338 | if not code.strip(): |
| 339 | raise typer.Exit(0) |
| 340 | |
| 341 | s.last_execution = ( |
| 342 | "stdin", |
| 343 | None, |
| 344 | datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), |
| 345 | ) |
| 346 | s.running = "repl(stdin)" |
| 347 | state.store.add(s) |
| 348 | try: |
| 349 | outputs = runtime.execute_code( |
| 350 | code, output_hook=lambda o: display_output(o, output_image) |
| 351 | ) |
| 352 | state.history.log_event( |
| 353 | name, "execution", {"code": code, "outputs": outputs, "source": "piped"} |
| 354 | ) |
| 355 | finally: |
| 356 | s.running = None |
| 357 | state.store.add(s) |
| 358 | runtime.stop() |
| 359 | else: |
| 360 | from colab_cli.repl import ColabREPL |
| 361 | |
| 362 | s.running = "repl" |
| 363 | state.store.add(s) |
| 364 | try: |
| 365 | repl_inst = ColabREPL( |
| 366 | runtime, |
| 367 | session_name=s.name, |
| 368 | history_logger=state.history, |
| 369 | output_image=output_image, |
| 370 | ) |
| 371 | state.history.log_event(name, "repl_started", {}) |
| 372 | repl_inst.run() |
| 373 | finally: |
| 374 | s.running = None |
| 375 | state.store.add(s) |
| 376 | |
| 377 | |
| 378 | def console( |
| 379 | session: Annotated[ |
| 380 | Optional[str], typer.Option("-s", "--session", help="Session name") |
| 381 | ] = None, |
| 382 | ): |
| 383 | """Connect to raw TTY console""" |
| 384 | from colab_cli.common import state |
| 385 | |
| 386 | name = state.resolve_session(session) |
| 387 | s = state.store.get(name) |
| 388 | if not s: |
| 389 | typer.echo(f"[colab] Session '{name}' not found.") |
| 390 | raise typer.Exit(1) |
| 391 | state.history.log_event(s.name, "console_started", {}) |
| 392 | s.running = "console" |
| 393 | state.store.add(s) |
| 394 | try: |
| 395 | connect_console(s) |
| 396 | except Exception as e: |
| 397 | if is_terminal_error(e): |
| 398 | typer.echo( |
| 399 | f"[colab] Session '{name}' appears to be lost (404/401). Cleaning up." |
| 400 | ) |
| 401 | state.prune_session(name) |
| 402 | raise typer.Exit(1) |
| 403 | raise e |
| 404 | finally: |
| 405 | s.running = None |
| 406 | state.store.add(s) |
| 407 | |
| 408 | |
| 409 | def register(app: typer.Typer): |
| 410 | app.command(name="exec")(exec_command) |
| 411 | app.command()(repl) |
| 412 | app.command()(console) |