google-colab-cli 0.5.4__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- colab_cli/auth.py +170 -0
- colab_cli/auto_update.py +226 -0
- colab_cli/cli.py +151 -0
- colab_cli/client.py +324 -0
- colab_cli/commands/__init__.py +14 -0
- colab_cli/commands/automation.py +265 -0
- colab_cli/commands/execution.py +356 -0
- colab_cli/commands/files.py +204 -0
- colab_cli/commands/run.py +471 -0
- colab_cli/commands/session.py +508 -0
- colab_cli/commands/utility.py +365 -0
- colab_cli/common.py +185 -0
- colab_cli/console.py +172 -0
- colab_cli/contents.py +93 -0
- colab_cli/converter.py +184 -0
- colab_cli/history.py +65 -0
- colab_cli/repl.py +173 -0
- colab_cli/runtime.py +262 -0
- colab_cli/state.py +152 -0
- colab_cli/utils.py +85 -0
- google_colab_cli-0.5.4.dist-info/METADATA +117 -0
- google_colab_cli-0.5.4.dist-info/RECORD +25 -0
- google_colab_cli-0.5.4.dist-info/WHEEL +4 -0
- google_colab_cli-0.5.4.dist-info/entry_points.txt +2 -0
- google_colab_cli-0.5.4.dist-info/licenses/LICENSE +202 -0
colab_cli/runtime.py
ADDED
|
@@ -0,0 +1,262 @@
|
|
|
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 logging
|
|
16
|
+
import time
|
|
17
|
+
from typing import Any, Callable, Dict, List, Optional
|
|
18
|
+
|
|
19
|
+
import jupyter_kernel_client
|
|
20
|
+
import requests
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class ColabRuntime:
|
|
24
|
+
def __init__(
|
|
25
|
+
self,
|
|
26
|
+
url: str,
|
|
27
|
+
token: str,
|
|
28
|
+
session_name: Optional[str] = None,
|
|
29
|
+
history: Optional[Any] = None,
|
|
30
|
+
kernel_id: Optional[str] = None,
|
|
31
|
+
session_id: Optional[str] = None,
|
|
32
|
+
on_kernel_started: Optional[Callable[[str], None]] = None,
|
|
33
|
+
on_session_started: Optional[Callable[[str], None]] = None,
|
|
34
|
+
):
|
|
35
|
+
self.url = url
|
|
36
|
+
self.token = token
|
|
37
|
+
self.session_name = session_name
|
|
38
|
+
self.history = history
|
|
39
|
+
self.kernel_id = kernel_id
|
|
40
|
+
self.session_id = session_id
|
|
41
|
+
self.on_kernel_started = on_kernel_started
|
|
42
|
+
self.on_session_started = on_session_started
|
|
43
|
+
self._kernel_client = None
|
|
44
|
+
self.colab_request_hook: Optional[Callable[[Dict[str, Any], Any], None]] = None
|
|
45
|
+
|
|
46
|
+
def _apply_ws_hook(self):
|
|
47
|
+
wsclient = self._kernel_client._manager.client
|
|
48
|
+
original_on_message = wsclient.kernel_socket.on_message
|
|
49
|
+
|
|
50
|
+
def hooked_on_message(s_ws, message):
|
|
51
|
+
if not self.colab_request_hook:
|
|
52
|
+
return original_on_message(s_ws, message)
|
|
53
|
+
|
|
54
|
+
try:
|
|
55
|
+
from jupyter_kernel_client.wsclient import JupyterSubprotocol
|
|
56
|
+
|
|
57
|
+
if wsclient._subprotocol == JupyterSubprotocol.DEFAULT:
|
|
58
|
+
from jupyter_kernel_client.wsclient import (
|
|
59
|
+
deserialize_msg_from_ws_default,
|
|
60
|
+
)
|
|
61
|
+
|
|
62
|
+
deserialize_msg = deserialize_msg_from_ws_default(message)
|
|
63
|
+
elif wsclient._subprotocol == JupyterSubprotocol.V1:
|
|
64
|
+
from jupyter_kernel_client.wsclient import (
|
|
65
|
+
deserialize_msg_from_ws_v1,
|
|
66
|
+
)
|
|
67
|
+
|
|
68
|
+
channel, msg_list = deserialize_msg_from_ws_v1(message)
|
|
69
|
+
deserialize_msg = wsclient.session.deserialize(msg_list)
|
|
70
|
+
else:
|
|
71
|
+
deserialize_msg = None
|
|
72
|
+
|
|
73
|
+
if deserialize_msg:
|
|
74
|
+
msg_type = deserialize_msg.get("msg_type")
|
|
75
|
+
if msg_type == "colab_request":
|
|
76
|
+
# We pass the deserialized msg and the wsclient to the hook
|
|
77
|
+
if self.colab_request_hook(deserialize_msg, wsclient):
|
|
78
|
+
# If the hook returns True, we intercept and do NOT pass to original
|
|
79
|
+
return
|
|
80
|
+
|
|
81
|
+
except Exception as e:
|
|
82
|
+
logging.debug(f"Error in colab_request hook: {e}")
|
|
83
|
+
|
|
84
|
+
# Call original for all other messages
|
|
85
|
+
original_on_message(s_ws, message)
|
|
86
|
+
|
|
87
|
+
wsclient.kernel_socket.on_message = hooked_on_message
|
|
88
|
+
|
|
89
|
+
@property
|
|
90
|
+
def kernel_client(self):
|
|
91
|
+
if not self._kernel_client:
|
|
92
|
+
retries = 3
|
|
93
|
+
backoff = 2
|
|
94
|
+
last_err = None
|
|
95
|
+
|
|
96
|
+
for i in range(retries):
|
|
97
|
+
try:
|
|
98
|
+
client_kwargs = {
|
|
99
|
+
"subprotocol": jupyter_kernel_client.JupyterSubprotocol.DEFAULT,
|
|
100
|
+
"extra_params": {"colab-runtime-proxy-token": self.token},
|
|
101
|
+
}
|
|
102
|
+
if self.session_id:
|
|
103
|
+
# WSSession (Session) expects 'session' for the ID
|
|
104
|
+
client_kwargs["session"] = self.session_id
|
|
105
|
+
|
|
106
|
+
self._kernel_client = jupyter_kernel_client.KernelClient(
|
|
107
|
+
server_url=self.url,
|
|
108
|
+
token=self.token,
|
|
109
|
+
kernel_id=self.kernel_id,
|
|
110
|
+
client_kwargs=client_kwargs,
|
|
111
|
+
headers={
|
|
112
|
+
"X-Colab-Client-Agent": "colab-cli",
|
|
113
|
+
"X-Colab-Runtime-Proxy-Token": self.token,
|
|
114
|
+
},
|
|
115
|
+
)
|
|
116
|
+
# Force _own_kernel to False. This prevents jupyter-kernel-client
|
|
117
|
+
# from automatically deleting the kernel when the client is closed or deleted.
|
|
118
|
+
self._kernel_client._own_kernel = False
|
|
119
|
+
|
|
120
|
+
self._kernel_client.start()
|
|
121
|
+
self._apply_ws_hook()
|
|
122
|
+
|
|
123
|
+
# Capture IDs if we started fresh
|
|
124
|
+
if not self.kernel_id and self._kernel_client.id:
|
|
125
|
+
self.kernel_id = self._kernel_client.id
|
|
126
|
+
if self.on_kernel_started:
|
|
127
|
+
self.on_kernel_started(self.kernel_id)
|
|
128
|
+
|
|
129
|
+
if (
|
|
130
|
+
not self.session_id
|
|
131
|
+
and self._kernel_client._manager.client.session.session
|
|
132
|
+
):
|
|
133
|
+
self.session_id = (
|
|
134
|
+
self._kernel_client._manager.client.session.session
|
|
135
|
+
)
|
|
136
|
+
if self.on_session_started:
|
|
137
|
+
self.on_session_started(self.session_id)
|
|
138
|
+
break
|
|
139
|
+
except (
|
|
140
|
+
requests.exceptions.ReadTimeout,
|
|
141
|
+
requests.exceptions.ConnectTimeout,
|
|
142
|
+
) as e:
|
|
143
|
+
last_err = e
|
|
144
|
+
if i < retries - 1:
|
|
145
|
+
sleep_time = backoff ** (i + 1)
|
|
146
|
+
logging.debug(
|
|
147
|
+
f"Kernel startup timeout, retrying in {sleep_time}s..."
|
|
148
|
+
f" ({i + 1}/{retries})"
|
|
149
|
+
)
|
|
150
|
+
time.sleep(sleep_time)
|
|
151
|
+
else:
|
|
152
|
+
raise last_err
|
|
153
|
+
except Exception as e:
|
|
154
|
+
raise e
|
|
155
|
+
|
|
156
|
+
return self._kernel_client
|
|
157
|
+
|
|
158
|
+
def restart(
|
|
159
|
+
self,
|
|
160
|
+
timeout: Optional[float] = None,
|
|
161
|
+
):
|
|
162
|
+
self.kernel_client.restart(timeout=timeout)
|
|
163
|
+
|
|
164
|
+
def execute_code(
|
|
165
|
+
self,
|
|
166
|
+
code: str,
|
|
167
|
+
allow_stdin: bool = False,
|
|
168
|
+
stdin_hook: Any = None,
|
|
169
|
+
output_hook: Optional[Callable[[Dict[str, Any]], None]] = None,
|
|
170
|
+
timeout: Optional[float] = None,
|
|
171
|
+
) -> List[Dict[str, Any]]:
|
|
172
|
+
# ``jupyter_kernel_client`` defaults ``timeout`` to ``REQUEST_TIMEOUT``
|
|
173
|
+
# (10 seconds) on both ``execute`` and ``execute_interactive``. That
|
|
174
|
+
# value is a wall-clock budget that shrinks every time the poll loop
|
|
175
|
+
# iterates -- as long as iopub/stdin events arrive back-to-back the
|
|
176
|
+
# call survives, but a single >10s quiet stretch (e.g. a kernel
|
|
177
|
+
# blocked on ``input_request`` while the user OAuths in the browser)
|
|
178
|
+
# will raise ``TimeoutError`` even though the underlying execution is
|
|
179
|
+
# still healthy. Callers that know they need a longer ceiling can
|
|
180
|
+
# pass ``timeout=`` here; otherwise we forward whatever the upstream
|
|
181
|
+
# default is (currently 10s).
|
|
182
|
+
kwargs = {"allow_stdin": allow_stdin}
|
|
183
|
+
if timeout is not None:
|
|
184
|
+
kwargs["timeout"] = timeout
|
|
185
|
+
|
|
186
|
+
# Wrap stdin_hook to log inputs
|
|
187
|
+
original_stdin_hook = stdin_hook
|
|
188
|
+
|
|
189
|
+
def wrapped_stdin_hook(prompt):
|
|
190
|
+
if self.history and self.session_name:
|
|
191
|
+
self.history.log_event(
|
|
192
|
+
self.session_name, "stdin_request", {"prompt": prompt}
|
|
193
|
+
)
|
|
194
|
+
|
|
195
|
+
res = original_stdin_hook(prompt) if original_stdin_hook else input(prompt)
|
|
196
|
+
|
|
197
|
+
if self.history and self.session_name:
|
|
198
|
+
self.history.log_event(self.session_name, "input_reply", {"value": res})
|
|
199
|
+
return res
|
|
200
|
+
|
|
201
|
+
if allow_stdin:
|
|
202
|
+
kwargs["stdin_hook"] = wrapped_stdin_hook
|
|
203
|
+
|
|
204
|
+
if output_hook:
|
|
205
|
+
# If we have an output hook, we use execute_interactive and manage buffering ourselves
|
|
206
|
+
outputs = []
|
|
207
|
+
|
|
208
|
+
def wrapped_output_hook(msg):
|
|
209
|
+
from jupyter_kernel_client.client import (
|
|
210
|
+
output_hook as default_output_hook,
|
|
211
|
+
)
|
|
212
|
+
|
|
213
|
+
# Update local outputs list using the default logic
|
|
214
|
+
new_indexes = default_output_hook(outputs, msg)
|
|
215
|
+
# If new outputs were added, call our streaming hook with the new data
|
|
216
|
+
if new_indexes:
|
|
217
|
+
for idx in sorted(new_indexes):
|
|
218
|
+
if idx < len(outputs):
|
|
219
|
+
output_hook(outputs[idx])
|
|
220
|
+
|
|
221
|
+
reply = self.kernel_client.execute_interactive(
|
|
222
|
+
code, output_hook=wrapped_output_hook, **kwargs
|
|
223
|
+
)
|
|
224
|
+
# execute_interactive returns the raw reply message
|
|
225
|
+
reply_content = reply["content"] if reply else {"status": "error"}
|
|
226
|
+
else:
|
|
227
|
+
reply = self.kernel_client.execute(code, **kwargs)
|
|
228
|
+
if not reply:
|
|
229
|
+
return []
|
|
230
|
+
outputs = reply.get("outputs", [])
|
|
231
|
+
reply_content = reply
|
|
232
|
+
|
|
233
|
+
# If there's an error status but no error in outputs, synthesize one
|
|
234
|
+
if reply_content.get("status") == "error":
|
|
235
|
+
has_error_output = any(o.get("output_type") == "error" for o in outputs)
|
|
236
|
+
if not has_error_output:
|
|
237
|
+
outputs.append(
|
|
238
|
+
{
|
|
239
|
+
"output_type": "error",
|
|
240
|
+
"ename": reply_content.get("ename", "Error"),
|
|
241
|
+
"evalue": reply_content.get("evalue", "Unknown error"),
|
|
242
|
+
"traceback": reply_content.get("traceback", []),
|
|
243
|
+
}
|
|
244
|
+
)
|
|
245
|
+
|
|
246
|
+
return outputs
|
|
247
|
+
|
|
248
|
+
def stop(self, shutdown_kernel: bool = False):
|
|
249
|
+
if self._kernel_client:
|
|
250
|
+
try:
|
|
251
|
+
# We manage kernel lifecycle explicitly.
|
|
252
|
+
# To prevent automatic shutdown, we bypass the manager's stop() and
|
|
253
|
+
# directly close the channels and socket.
|
|
254
|
+
client = self._kernel_client._manager.client
|
|
255
|
+
client.stop_channels()
|
|
256
|
+
if client.kernel_socket:
|
|
257
|
+
client.kernel_socket.close()
|
|
258
|
+
|
|
259
|
+
if shutdown_kernel:
|
|
260
|
+
self._kernel_client._manager.shutdown_kernel(now=True)
|
|
261
|
+
except Exception:
|
|
262
|
+
logging.exception("Error stopping kernel client")
|
colab_cli/state.py
ADDED
|
@@ -0,0 +1,152 @@
|
|
|
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 contextlib
|
|
16
|
+
import json
|
|
17
|
+
import os
|
|
18
|
+
import fcntl
|
|
19
|
+
from datetime import datetime
|
|
20
|
+
from typing import Dict, Optional, Tuple, Iterator, IO
|
|
21
|
+
from pydantic import BaseModel
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
class SessionState(BaseModel):
|
|
25
|
+
name: str
|
|
26
|
+
token: str
|
|
27
|
+
url: str
|
|
28
|
+
endpoint: str
|
|
29
|
+
variant: str = "DEFAULT"
|
|
30
|
+
accelerator: str = "NONE"
|
|
31
|
+
kernel_id: Optional[str] = None
|
|
32
|
+
session_id: Optional[str] = None
|
|
33
|
+
last_execution: Optional[Tuple[str, Optional[str], str]] = None
|
|
34
|
+
running: Optional[str] = None
|
|
35
|
+
keep_alive_pid: Optional[int] = None
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
class Settings(BaseModel):
|
|
39
|
+
update_url: str = "https://pypi.org/pypi/google-colab-cli/json"
|
|
40
|
+
last_check: Optional[datetime] = None
|
|
41
|
+
enable_update_check: bool = True
|
|
42
|
+
# Highest version seen on the update source; cached for the banner.
|
|
43
|
+
latest_version: Optional[str] = None
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
class _LockedFileStore:
|
|
47
|
+
def __init__(self, path: str):
|
|
48
|
+
self.path = path
|
|
49
|
+
self._ensure_dir()
|
|
50
|
+
|
|
51
|
+
def _ensure_dir(self):
|
|
52
|
+
os.makedirs(os.path.dirname(self.path), exist_ok=True)
|
|
53
|
+
|
|
54
|
+
def _write_data(self, f: IO, data: str):
|
|
55
|
+
f.seek(0)
|
|
56
|
+
f.truncate()
|
|
57
|
+
f.write(data)
|
|
58
|
+
f.flush()
|
|
59
|
+
os.fsync(f.fileno())
|
|
60
|
+
|
|
61
|
+
@contextlib.contextmanager
|
|
62
|
+
def _lock_shared(self) -> Iterator[Optional[IO]]:
|
|
63
|
+
if not os.path.exists(self.path):
|
|
64
|
+
yield None
|
|
65
|
+
return
|
|
66
|
+
with open(self.path, "r") as f:
|
|
67
|
+
fcntl.flock(f, fcntl.LOCK_SH)
|
|
68
|
+
try:
|
|
69
|
+
yield f
|
|
70
|
+
finally:
|
|
71
|
+
fcntl.flock(f, fcntl.LOCK_UN)
|
|
72
|
+
|
|
73
|
+
@contextlib.contextmanager
|
|
74
|
+
def _lock_exclusive(self) -> Iterator[IO]:
|
|
75
|
+
with open(self.path, "a+") as f:
|
|
76
|
+
fcntl.flock(f, fcntl.LOCK_EX)
|
|
77
|
+
try:
|
|
78
|
+
yield f
|
|
79
|
+
finally:
|
|
80
|
+
fcntl.flock(f, fcntl.LOCK_UN)
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
class SettingsStore(_LockedFileStore):
|
|
84
|
+
def __init__(self, path: Optional[str] = None):
|
|
85
|
+
if not path:
|
|
86
|
+
path = os.path.expanduser("~/.config/colab-cli/settings.json")
|
|
87
|
+
super().__init__(path)
|
|
88
|
+
|
|
89
|
+
def load(self) -> Settings:
|
|
90
|
+
with self._lock_shared() as f:
|
|
91
|
+
if f is None:
|
|
92
|
+
return Settings()
|
|
93
|
+
try:
|
|
94
|
+
content = f.read()
|
|
95
|
+
if not content or content.isspace():
|
|
96
|
+
return Settings()
|
|
97
|
+
data = json.loads(content)
|
|
98
|
+
return Settings.model_validate(data)
|
|
99
|
+
except Exception:
|
|
100
|
+
return Settings()
|
|
101
|
+
|
|
102
|
+
def save(self, settings: Settings):
|
|
103
|
+
with self._lock_exclusive() as f:
|
|
104
|
+
self._write_data(f, settings.model_dump_json(indent=2))
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
class StateStore(_LockedFileStore):
|
|
108
|
+
def __init__(self, path: Optional[str] = None):
|
|
109
|
+
if not path:
|
|
110
|
+
path = os.path.expanduser("~/.config/colab-cli/sessions.json")
|
|
111
|
+
super().__init__(path)
|
|
112
|
+
|
|
113
|
+
def _load_raw(self, f) -> Dict[str, SessionState]:
|
|
114
|
+
try:
|
|
115
|
+
f.seek(0)
|
|
116
|
+
content = f.read()
|
|
117
|
+
if not content or content.isspace():
|
|
118
|
+
return {}
|
|
119
|
+
data = json.loads(content)
|
|
120
|
+
return {k: SessionState(**v) for k, v in data.items()}
|
|
121
|
+
except Exception:
|
|
122
|
+
return {}
|
|
123
|
+
|
|
124
|
+
def _save_raw(self, f, sessions: Dict[str, SessionState]):
|
|
125
|
+
content = json.dumps({k: v.model_dump() for k, v in sessions.items()}, indent=2)
|
|
126
|
+
self._write_data(f, content)
|
|
127
|
+
|
|
128
|
+
def add(self, state: SessionState):
|
|
129
|
+
with self._lock_exclusive() as f:
|
|
130
|
+
sessions = self._load_raw(f)
|
|
131
|
+
sessions[state.name] = state
|
|
132
|
+
self._save_raw(f, sessions)
|
|
133
|
+
|
|
134
|
+
def get(self, name: str) -> Optional[SessionState]:
|
|
135
|
+
with self._lock_shared() as f:
|
|
136
|
+
if f is None:
|
|
137
|
+
return None
|
|
138
|
+
sessions = self._load_raw(f)
|
|
139
|
+
return sessions.get(name)
|
|
140
|
+
|
|
141
|
+
def remove(self, name: str):
|
|
142
|
+
with self._lock_exclusive() as f:
|
|
143
|
+
sessions = self._load_raw(f)
|
|
144
|
+
if name in sessions:
|
|
145
|
+
del sessions[name]
|
|
146
|
+
self._save_raw(f, sessions)
|
|
147
|
+
|
|
148
|
+
def list(self) -> Dict[str, SessionState]:
|
|
149
|
+
with self._lock_shared() as f:
|
|
150
|
+
if f is None:
|
|
151
|
+
return {}
|
|
152
|
+
return self._load_raw(f)
|
colab_cli/utils.py
ADDED
|
@@ -0,0 +1,85 @@
|
|
|
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 base64
|
|
16
|
+
import logging
|
|
17
|
+
import sys
|
|
18
|
+
import tempfile
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
from typing import Optional
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def get_status_code(e: Exception) -> Optional[int]:
|
|
25
|
+
"""Safely extracts status code from various exception types."""
|
|
26
|
+
if hasattr(e, "response") and e.response is not None:
|
|
27
|
+
if hasattr(e.response, "status_code"):
|
|
28
|
+
return e.response.status_code
|
|
29
|
+
if hasattr(e, "status_code"):
|
|
30
|
+
return e.status_code
|
|
31
|
+
return None
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def is_terminal_error(e: Exception) -> bool:
|
|
35
|
+
"""Checks if an exception indicates a lost session (404/401)."""
|
|
36
|
+
code = get_status_code(e)
|
|
37
|
+
if code in (404, 401):
|
|
38
|
+
return True
|
|
39
|
+
# Some exceptions from jupyter-kernel-client might wrap the real one or be different
|
|
40
|
+
err_msg = str(e)
|
|
41
|
+
if "404" in err_msg or "401" in err_msg:
|
|
42
|
+
return True
|
|
43
|
+
return False
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def print_kitty(image_bytes: bytes):
|
|
47
|
+
"""
|
|
48
|
+
Outputs an image using the Kitty Graphics Protocol.
|
|
49
|
+
Expects PNG bytes.
|
|
50
|
+
|
|
51
|
+
No-op when stdout is not a TTY: the escape sequence is meaningless to a
|
|
52
|
+
file/pipe and visually corrupts captured output (e.g. when piping
|
|
53
|
+
`colab exec` into a shell tool, redirecting to a log file, or running
|
|
54
|
+
under non-Kitty terminals). Callers still get the image via
|
|
55
|
+
`handle_image`'s file-write path.
|
|
56
|
+
"""
|
|
57
|
+
if not sys.stdout.isatty():
|
|
58
|
+
return
|
|
59
|
+
try:
|
|
60
|
+
b64_data = base64.b64encode(image_bytes).decode("ascii")
|
|
61
|
+
sys.stdout.write("\n\033_Ga=T,f=100;")
|
|
62
|
+
sys.stdout.write(b64_data)
|
|
63
|
+
sys.stdout.write("\033\\\n")
|
|
64
|
+
sys.stdout.flush()
|
|
65
|
+
except Exception:
|
|
66
|
+
logging.exception("Kitty rendering failed")
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def handle_image(image_b64: str, mime_type: str = "image/png", target_path: str = None):
|
|
70
|
+
image_bytes = base64.b64decode(image_b64)
|
|
71
|
+
# Print inline using Kitty protocol
|
|
72
|
+
print_kitty(image_bytes)
|
|
73
|
+
|
|
74
|
+
if target_path:
|
|
75
|
+
# If a target path is specified, save it there
|
|
76
|
+
with open(target_path, "wb") as f:
|
|
77
|
+
f.write(image_bytes)
|
|
78
|
+
print(f"\n[Image saved to: {target_path}]")
|
|
79
|
+
else:
|
|
80
|
+
# Save to temp file as fallback
|
|
81
|
+
ext = mime_type.split("/")[-1]
|
|
82
|
+
tmp = tempfile.NamedTemporaryFile(delete=False, suffix=f".{ext}")
|
|
83
|
+
tmp.write(image_bytes)
|
|
84
|
+
tmp.close()
|
|
85
|
+
print(f"\n[Image saved to: {tmp.name}]")
|
|
@@ -0,0 +1,117 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: google-colab-cli
|
|
3
|
+
Version: 0.5.4
|
|
4
|
+
Summary: CLI for interacting with Colab.
|
|
5
|
+
License-File: LICENSE
|
|
6
|
+
Requires-Python: >=3.13
|
|
7
|
+
Requires-Dist: google-auth-oauthlib>=1.3.0
|
|
8
|
+
Requires-Dist: google-auth>=2.49.1
|
|
9
|
+
Requires-Dist: jupyter-kernel-client
|
|
10
|
+
Requires-Dist: nbformat>=5.10.4
|
|
11
|
+
Requires-Dist: packaging>=24.0
|
|
12
|
+
Requires-Dist: prompt-toolkit>=3.0.52
|
|
13
|
+
Requires-Dist: pydantic>=2.12.5
|
|
14
|
+
Requires-Dist: pygments>=2.19.2
|
|
15
|
+
Requires-Dist: pytest-cov>=7.0.0
|
|
16
|
+
Requires-Dist: pytest-mock>=3.15.1
|
|
17
|
+
Requires-Dist: pytest>=9.0.2
|
|
18
|
+
Requires-Dist: requests>=2.32.5
|
|
19
|
+
Requires-Dist: rich>=14.3.3
|
|
20
|
+
Requires-Dist: typer>=0.24.1
|
|
21
|
+
Description-Content-Type: text/markdown
|
|
22
|
+
|
|
23
|
+
# Colab CLI
|
|
24
|
+
|
|
25
|
+
A command-line interface for Google Colab. Create sessions, run code, manage files, and capture work — all without leaving your terminal.
|
|
26
|
+
|
|
27
|
+
> Why? The agents are coming.
|
|
28
|
+
|
|
29
|
+
## Install
|
|
30
|
+
|
|
31
|
+
```bash
|
|
32
|
+
uv tool install git+https://github.com/googlecolab/google-colab-cli
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
Note: If you have a non-standard default package index (Googlers), you may also need to add `--index https://pypi.org/simple`.
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
## Quick start
|
|
39
|
+
|
|
40
|
+
```bash
|
|
41
|
+
colab new # provision a CPU session
|
|
42
|
+
echo "print('hello')" | colab exec # run code
|
|
43
|
+
colab stop # release the VM
|
|
44
|
+
```
|
|
45
|
+
|
|
46
|
+
When only one session is active you can omit `-s <session>`; the CLI selects it automatically.
|
|
47
|
+
|
|
48
|
+
## Commands
|
|
49
|
+
|
|
50
|
+
### Sessions
|
|
51
|
+
| Command | Description |
|
|
52
|
+
| --- | --- |
|
|
53
|
+
| `colab new [-s NAME] [--gpu T4\|L4\|A100\|H100] [--tpu v5e1\|v6e1]` | Provision a new session (CPU by default) |
|
|
54
|
+
| `colab sessions` | List all active sessions on the backend |
|
|
55
|
+
| `colab status [-s NAME]` | Show one session, or all locally-known sessions |
|
|
56
|
+
| `colab stop -s NAME` | Terminate a session |
|
|
57
|
+
| `colab url [-s NAME] [--open]` | Print a browser URL that opens the session in Colab |
|
|
58
|
+
|
|
59
|
+
### Execution
|
|
60
|
+
| Command | Description |
|
|
61
|
+
| --- | --- |
|
|
62
|
+
| `colab exec [-s NAME] [-f FILE] [--output-image PATH]` | Run Python from stdin, a `.py` file, or a `.ipynb` notebook |
|
|
63
|
+
| `colab repl [-s NAME] [--output-image PATH]` | Interactive Python REPL (or one-shot if stdin is piped) |
|
|
64
|
+
| `colab console [-s NAME]` | Raw TTY shell on the VM (or one-shot if stdin is piped) |
|
|
65
|
+
|
|
66
|
+
### Files
|
|
67
|
+
| Command | Description |
|
|
68
|
+
| --- | --- |
|
|
69
|
+
| `colab ls [-s NAME] [PATH]` | List remote files |
|
|
70
|
+
| `colab upload -s NAME LOCAL REMOTE` | Upload a file |
|
|
71
|
+
| `colab download -s NAME REMOTE LOCAL` | Download a file |
|
|
72
|
+
| `colab rm -s NAME PATH` | Delete a remote file |
|
|
73
|
+
| `colab edit -s NAME PATH` | Edit a remote file in `$EDITOR` |
|
|
74
|
+
|
|
75
|
+
### Automation & utility
|
|
76
|
+
| Command | Description |
|
|
77
|
+
| --- | --- |
|
|
78
|
+
| `colab auth -s NAME` | Authenticate the VM for GCP services |
|
|
79
|
+
| `colab drivemount -s NAME [PATH]` | Mount Google Drive (default `/content/drive`) |
|
|
80
|
+
| `colab install -s NAME [-r requirements.txt \| pkg ...]` | Install packages with `uv` (falls back to `pip`) |
|
|
81
|
+
| `colab log [-s NAME] [-n N] [-o FILE]` | View or export session history (`.ipynb`/`.md`/`.txt`/`.jsonl`) |
|
|
82
|
+
| `colab pay` | Open the Colab signup page |
|
|
83
|
+
| `colab version` | Print the installed version |
|
|
84
|
+
| `colab update [--install]` | Check for a newer release (and optionally install it) |
|
|
85
|
+
| `colab help` | Show usage |
|
|
86
|
+
|
|
87
|
+
### Global options
|
|
88
|
+
- `--auth {oauth2,adc}` — authentication strategy (default `oauth2`)
|
|
89
|
+
- `-c, --client-oauth-config PATH` — OAuth client config (default `~/.colab-cli-oauth-config.json`)
|
|
90
|
+
- `--config PATH` — session state file (default `~/.config/colab-cli/sessions.json`)
|
|
91
|
+
- `--logtostderr` — send all output to stderr
|
|
92
|
+
|
|
93
|
+
## Examples
|
|
94
|
+
|
|
95
|
+
```bash
|
|
96
|
+
# Train a model on an A100, save the checkpoint locally
|
|
97
|
+
colab new -s trainer --gpu A100
|
|
98
|
+
colab install -s trainer torch transformers
|
|
99
|
+
colab exec -s trainer -f train.py
|
|
100
|
+
colab download -s trainer checkpoints/model.bin ./model.bin
|
|
101
|
+
colab stop -s trainer
|
|
102
|
+
|
|
103
|
+
# Mount Drive and analyze a notebook
|
|
104
|
+
colab new -s analysis
|
|
105
|
+
colab drivemount -s analysis
|
|
106
|
+
colab exec -s analysis -f analysis.ipynb # writes analysis_output.ipynb
|
|
107
|
+
colab log -s analysis -o report.ipynb
|
|
108
|
+
colab stop -s analysis
|
|
109
|
+
```
|
|
110
|
+
|
|
111
|
+
## Notes
|
|
112
|
+
|
|
113
|
+
- `repl` and `console` require a TTY when run interactively. Pipe stdin to use them in scripts.
|
|
114
|
+
- `exec` reads files locally and ships their contents to the VM — local edits don't require uploading.
|
|
115
|
+
- Session metadata is stored at `~/.config/colab-cli/sessions.json`. Settings (auto-update etc.) live at `~/.config/colab-cli/settings.json`.
|
|
116
|
+
|
|
117
|
+
See [`CONTRIBUTING.md`](./CONTRIBUTING.md) for how to file feedback.
|
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
colab_cli/auth.py,sha256=OGhq_thOkZkEttSgEAeVnOjIEqVtwkSVB2HeZz6f01I,6586
|
|
2
|
+
colab_cli/auto_update.py,sha256=JCQxgv8GwIlUPM36gQeZSk0pqhrJj7AVZ5fCZ2FP86k,7476
|
|
3
|
+
colab_cli/cli.py,sha256=b4iCzBKiUtI-q6yJ4kOlvjFOb0XJ02gnBmqPdcjVZtY,4481
|
|
4
|
+
colab_cli/client.py,sha256=mDIE4TgBgZAbCVCax_D9dp6CxeINk24QkhrMBIw50sE,10780
|
|
5
|
+
colab_cli/common.py,sha256=J3xNCv43Xkgq4eQdBJtz15r41_BYzXPjdvob3fj3nw0,6368
|
|
6
|
+
colab_cli/console.py,sha256=0snB6MQr-gEH7x_P87pPVyWBslSEJIZCtJuIx1S_LGA,5943
|
|
7
|
+
colab_cli/contents.py,sha256=6vQspKs6zws88blwmJcniUjfVfYLCOW5lheNQG_Gs7Y,3026
|
|
8
|
+
colab_cli/converter.py,sha256=IRFDy7o053DvEmyRBtCIMrgiH_Ol52j8-oUd2w1K3LI,7118
|
|
9
|
+
colab_cli/history.py,sha256=0fHvipAZXdddzBz7LBqKYDJq-gC5M3D28aqwQeHrrbQ,2332
|
|
10
|
+
colab_cli/repl.py,sha256=rZMU34opVw-gUy-F9oK2NQEs7hlPz8cu4vtbRjTdv-4,5697
|
|
11
|
+
colab_cli/runtime.py,sha256=njZZ8zOnUkFvnGdVTIxZVOT552NPSG54H79uZQKG5v8,10734
|
|
12
|
+
colab_cli/state.py,sha256=E6LueTijt6skatAbs8GLMdU68JDI5gFFvlXTwlailH0,4792
|
|
13
|
+
colab_cli/utils.py,sha256=tuWWmM3GJdg3yUnv1lvvSLg6Pcox5jnjjYrqcIy3uBc,2865
|
|
14
|
+
colab_cli/commands/__init__.py,sha256=A2GIGT_cfHV0jkPs0o9RPx0pHZP459CQoRn8o8KkR0E,575
|
|
15
|
+
colab_cli/commands/automation.py,sha256=bmVq_x4EVtkkuWmp1tha-Ilag9OwX1XX9YLpZda0bRo,9290
|
|
16
|
+
colab_cli/commands/execution.py,sha256=GbXmM_dYQayPIA9LSAiVrlRL8ivQulaoZfhjAcB070s,11222
|
|
17
|
+
colab_cli/commands/files.py,sha256=VvguO3ftcPdl74HznvoKCmQw52257Fl64IqlWtFd7z4,6503
|
|
18
|
+
colab_cli/commands/run.py,sha256=21fnKP5ceK6wXT6bM2w4au6esh7nIgWyIDismGVoS7c,15674
|
|
19
|
+
colab_cli/commands/session.py,sha256=FycsSctSWlv-NpY9-_mI_zguE6j_SNEmVjL0edS2pvE,17016
|
|
20
|
+
colab_cli/commands/utility.py,sha256=5HyicD91XEhpOBprt_9csNpocvNzNNHZZLezjNG2amI,13369
|
|
21
|
+
google_colab_cli-0.5.4.dist-info/METADATA,sha256=2QwHjj4sjNlPW79buTBvRKwH8GW3rmaR-VvzxDXbgLI,4424
|
|
22
|
+
google_colab_cli-0.5.4.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
|
|
23
|
+
google_colab_cli-0.5.4.dist-info/entry_points.txt,sha256=zepj_5XZElFi_sSsIf6_9i5JDbPFj71nFwvGFJU4mxc,45
|
|
24
|
+
google_colab_cli-0.5.4.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
|
|
25
|
+
google_colab_cli-0.5.4.dist-info/RECORD,,
|