| 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 | from typing import Any, List, Optional |
| 17 | |
| 18 | from prompt_toolkit import PromptSession |
| 19 | from prompt_toolkit.history import InMemoryHistory |
| 20 | from prompt_toolkit.key_binding import KeyBindings |
| 21 | from prompt_toolkit.lexers import PygmentsLexer |
| 22 | from prompt_toolkit.styles import Style |
| 23 | from pygments.lexers.python import PythonLexer |
| 24 | from rich.console import Console |
| 25 | from rich.text import Text |
| 26 | |
| 27 | from colab_cli.runtime import ColabRuntime |
| 28 | from colab_cli.utils import handle_image, render_display_data |
| 29 | |
| 30 | |
| 31 | |
| 32 | class ColabREPL: |
| 33 | def __init__( |
| 34 | self, |
| 35 | runtime: ColabRuntime, |
| 36 | session_name: Optional[str] = None, |
| 37 | history_logger: Optional[Any] = None, |
| 38 | output_image: Optional[str] = None, |
| 39 | ): |
| 40 | self.runtime = runtime |
| 41 | self.session_name = session_name |
| 42 | self.history_logger = history_logger |
| 43 | self.output_image = output_image |
| 44 | self.kb = KeyBindings() |
| 45 | self.console = Console() |
| 46 | self.repl_history: List[dict] = [] |
| 47 | |
| 48 | @self.kb.add("enter") |
| 49 | def _(event): |
| 50 | event.current_buffer.validate_and_handle() |
| 51 | |
| 52 | @self.kb.add("escape", "enter") |
| 53 | @self.kb.add("c-j") |
| 54 | def _(event): |
| 55 | event.current_buffer.insert_text("\n") |
| 56 | |
| 57 | self.session = PromptSession( |
| 58 | history=InMemoryHistory(), |
| 59 | lexer=PygmentsLexer(PythonLexer), |
| 60 | include_default_pygments_style=False, |
| 61 | key_bindings=self.kb, |
| 62 | multiline=True, |
| 63 | ) |
| 64 | self.style = Style.from_dict( |
| 65 | { |
| 66 | "prompt": "bold blue", |
| 67 | "continuation": "#888888", |
| 68 | } |
| 69 | ) |
| 70 | |
| 71 | def print_info(self, message: str): |
| 72 | self.console.print(f"[bold blue][*][/bold blue] {message}") |
| 73 | |
| 74 | def print_error(self, message: str): |
| 75 | self.console.print(f"[bold red][!][/bold red] {message}") |
| 76 | |
| 77 | def display_output(self, output: dict): |
| 78 | if "text" in output: |
| 79 | self.console.print(Text.from_ansi(output["text"]), end="") |
| 80 | elif "data" in output: |
| 81 | data = output["data"] |
| 82 | |
| 83 | # Check for images first |
| 84 | image_displayed = False |
| 85 | for mime_type in ["image/png", "image/jpeg"]: |
| 86 | if mime_type in data: |
| 87 | handle_image( |
| 88 | data[mime_type], mime_type, target_path=self.output_image |
| 89 | ) |
| 90 | image_displayed = True |
| 91 | break |
| 92 | |
| 93 | text = render_display_data(data) |
| 94 | if text is not None: |
| 95 | # Skip generic IPython object reprs if we already showed an image |
| 96 | if isinstance(text, Text) and image_displayed: |
| 97 | if any( |
| 98 | x in text.plain |
| 99 | for x in ["<IPython.core.display.Image", "<Figure size"] |
| 100 | ): |
| 101 | return |
| 102 | self.console.print(text) |
| 103 | elif output.get("output_type") == "error": |
| 104 | ename = output.get("ename", "Error") |
| 105 | evalue = output.get("evalue", "") |
| 106 | traceback = output.get("traceback", []) |
| 107 | if traceback: |
| 108 | self.console.print(Text.from_ansi("".join(traceback))) |
| 109 | else: |
| 110 | self.print_error(f"{ename}: {evalue}") |
| 111 | |
| 112 | def execute(self, code: str): |
| 113 | if self.session_name: |
| 114 | from colab_cli.common import state |
| 115 | |
| 116 | s = state.store.get(self.session_name) |
| 117 | if s: |
| 118 | s.last_execution = ( |
| 119 | "REPL", |
| 120 | None, |
| 121 | datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), |
| 122 | ) |
| 123 | state.store.add(s) |
| 124 | |
| 125 | try: |
| 126 | outputs = self.runtime.execute_code( |
| 127 | code, output_hook=lambda o: self.display_output(o) |
| 128 | ) |
| 129 | # Ensure next prompt starts on a newline after streaming |
| 130 | print() |
| 131 | |
| 132 | self.repl_history.append({"input": code, "outputs": outputs or []}) |
| 133 | if self.history_logger and self.session_name: |
| 134 | self.history_logger.log_event( |
| 135 | self.session_name, |
| 136 | "execution", |
| 137 | {"code": code, "outputs": outputs or []}, |
| 138 | ) |
| 139 | except Exception as e: |
| 140 | self.print_error(f"Execution failed: {e}") |
| 141 | |
| 142 | def run(self): |
| 143 | self.console.print("Python 3 (Google Colab Runtime)\nType /quit to exit.") |
| 144 | |
| 145 | while True: |
| 146 | try: |
| 147 | result = self.session.prompt( |
| 148 | ">>> ", |
| 149 | style=self.style, |
| 150 | ) |
| 151 | |
| 152 | if result is None: |
| 153 | continue |
| 154 | |
| 155 | code = result.strip() |
| 156 | |
| 157 | if not code: |
| 158 | continue |
| 159 | |
| 160 | if code.lower() in ("/quit", "quit()", "exit()"): |
| 161 | break |
| 162 | |
| 163 | self.execute(code) |
| 164 | |
| 165 | except EOFError: |
| 166 | break |
| 167 | except KeyboardInterrupt: |
| 168 | print() |
| 169 | continue |
| 170 | except Exception as e: |
| 171 | self.print_error(f"REPL Error: {e}") |
| 172 | |
| 173 | self.print_info("Goodbye!") |
| 174 | self.runtime.stop() |