feat: Add --env KEY=VALUE flag to colab run and colab exec to inject environment variables (#65)
Matt Van Horn committed
Jul 30, 2026 at 13:04 UTC
507d169533efd95662274a79309f568dddd4215f
4 files changed
+300
-5
src/colab_cli/commands/execution.py
+48
-3
@@ -21,7 +21,7 @@ import typer
21
import uuid
22
from nbformat.v4 import new_output
23
from rich.console import Console
24
-from typing import Optional
24
+from typing import List, Optional
25
from typing_extensions import Annotated
26
27
from colab_cli.runtime import ColabRuntime
@@ -31,12 +31,47 @@ from colab_cli.console import connect_console
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
@@ -75,7 +110,6 @@ def save_output(outputs, cell):
110
)
111
112
78
-
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
@@ -117,10 +151,21 @@ def exec_command(
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:
@@ -190,7 +235,7 @@ def exec_command(
235
state.store.add(s)
236
237
for i, block in enumerate(code_blocks):
193
- code = block["code"]
238
+ code = _build_env_prelude(env_vars) + block["code"]
239
identifier = None
240
if is_nb:
241
title_match = TITLE_REGEX.search(code)
src/colab_cli/commands/run.py
+50
-2
@@ -44,6 +44,7 @@ from colab_cli.client import (
44
PostAssignmentResponse,
45
Variant,
46
)
47
+from colab_cli.commands.execution import _build_env_prelude, _parse_env_vars
48
from colab_cli.commands.session import (
49
_is_scope_error,
50
_scope_remediation_message,
@@ -75,12 +76,15 @@ def _resolve_accelerator(gpu: Optional[str], tpu: Optional[str]):
76
return Variant.DEFAULT, Accelerator.NONE
77
78
78
-def _build_script_payload(script_path: str, script_args: List[str]) -> str:
79
+def _build_script_payload(
80
+ script_path: str, script_args: List[str], env_vars: Optional[dict[str, str]] = None
81
+) -> str:
82
"""Wrap the script body so it executes with native-`python`-like semantics.
83
84
Specifically:
85
- `sys.argv = [<basename>, *script_args]` so `argparse` etc. work.
86
- `__name__ = '__main__'` so `if __name__ == "__main__":` guards fire.
87
+ - Requested `--env KEY=VALUE` pairs are written into `os.environ`.
88
- Suppress the IPython UserWarning "To exit: use 'exit', 'quit', or
89
Ctrl-D." which fires whenever the script calls `sys.exit(...)`. This
90
warning is meaningful in an interactive REPL, but for `colab run` it
@@ -102,10 +106,42 @@ def _build_script_payload(script_path: str, script_args: List[str]) -> str:
106
f"sys.argv = {argv_literal}\n"
107
"__name__ = '__main__'\n"
108
"warnings.filterwarnings('ignore', message=\"To exit: use\")\n"
109
+ + _build_env_prelude(env_vars or {})
110
+ _strip_shebang(body)
111
)
112
113
114
+def _extract_env_args_from_script_args(
115
+ script_args: List[str],
116
+) -> tuple[List[str], List[str]]:
117
+ """Pull colab's --env options out of variadic script args.
118
+
119
+ `colab run` intentionally forwards unknown options after the script path to
120
+ the user's script. Since `--env` is now a colab option, support both
121
+ `colab run --env KEY=VALUE script.py` (Typer parses this) and
122
+ `colab run script.py --env KEY=VALUE` (this helper parses it).
123
+ """
124
+ forwarded_args = []
125
+ env = []
126
+ i = 0
127
+ while i < len(script_args):
128
+ arg = script_args[i]
129
+ if arg == "--env":
130
+ if i + 1 >= len(script_args):
131
+ env.append("")
132
+ i += 1
133
+ else:
134
+ env.append(script_args[i + 1])
135
+ i += 2
136
+ elif arg.startswith("--env="):
137
+ env.append(arg.split("=", 1)[1])
138
+ i += 1
139
+ else:
140
+ forwarded_args.append(arg)
141
+ i += 1
142
+ return forwarded_args, env
143
+
144
+
145
def _strip_shebang(body: str) -> str:
146
"""Remove a leading `#!...\\n` if present. The remote kernel doesn't need
147
or understand it (it's a contract between the local kernel and the file's
@@ -237,6 +273,16 @@ def run_command(
273
Optional[float],
274
typer.Option("--timeout", help="Timeout in seconds for code execution"),
275
] = 30.0,
276
+ env: Annotated[
277
+ Optional[List[str]],
278
+ typer.Option(
279
+ "--env",
280
+ help=(
281
+ "Set an environment variable in the remote kernel as KEY=VALUE. "
282
+ "Repeat for multiple variables."
283
+ ),
284
+ ),
285
+ ] = None,
286
):
287
"""Run a Python script on a fresh Colab VM, then release the VM
288
@@ -250,6 +296,8 @@ def run_command(
296
from colab_cli.common import state
297
298
script_args = script_args or []
299
+ script_args, inline_env = _extract_env_args_from_script_args(script_args)
300
+ env_vars = _parse_env_vars([*(env or []), *inline_env])
301
302
# AGENTS.md item 10: validate locally BEFORE allocating a VM. A typo'd
303
# script path should not cost the user real compute.
@@ -379,7 +427,7 @@ def run_command(
427
raise typer.Exit(1)
428
raise
429
382
- payload = _build_script_payload(script, script_args)
430
+ payload = _build_script_payload(script, script_args, env_vars)
431
s.running = f"run({os.path.basename(script)})"
432
s.last_execution = (
433
script,
tests/test_exec.py
+94
@@ -87,6 +87,100 @@ def test_cli_exec_stdin(mock_store, mock_runtime_class, mock_common_state):
87
)
88
89
90
+def test_cli_exec_env_injects_prelude(
91
+ mock_store, mock_runtime_class, mock_common_state
92
+):
93
+ mock_session = MagicMock()
94
+ mock_session.name = "s1"
95
+ mock_session.url = "http://url"
96
+ mock_session.token = "token"
97
+ mock_session.kernel_id = None
98
+ mock_session.session_id = None
99
+ mock_store.get.return_value = mock_session
100
+
101
+ mock_common_state.resolve_session.return_value = "s1"
102
+ mock_runtime = mock_runtime_class.return_value
103
+ mock_runtime.execute_code.return_value = []
104
+
105
+ code = "import os\nprint(os.environ.get('HF_TOKEN'))"
106
+ result = runner.invoke(
107
+ app, ["exec", "-s", "s1", "--env", "HF_TOKEN=abc"], input=code
108
+ )
109
+
110
+ assert result.exit_code == 0, result.output
111
+ expected = "import os\nos.environ['HF_TOKEN'] = 'abc'\n" + code
112
+ mock_runtime.execute_code.assert_any_call(expected, output_hook=ANY, timeout=30.0)
113
+
114
+
115
+def test_cli_exec_env_flags_accumulate_and_split_on_first_equals(
116
+ mock_store, mock_runtime_class, mock_common_state
117
+):
118
+ mock_session = MagicMock()
119
+ mock_session.name = "s1"
120
+ mock_session.url = "http://url"
121
+ mock_session.token = "token"
122
+ mock_session.kernel_id = None
123
+ mock_session.session_id = None
124
+ mock_store.get.return_value = mock_session
125
+
126
+ mock_common_state.resolve_session.return_value = "s1"
127
+ mock_runtime = mock_runtime_class.return_value
128
+ mock_runtime.execute_code.return_value = []
129
+
130
+ result = runner.invoke(
131
+ app,
132
+ ["exec", "-s", "s1", "--env", "HF_TOKEN=abc", "--env", "B64=a=b=c"],
133
+ input="print('ok')",
134
+ )
135
+
136
+ assert result.exit_code == 0, result.output
137
+ expected = (
138
+ "import os\n"
139
+ "os.environ['HF_TOKEN'] = 'abc'\n"
140
+ "os.environ['B64'] = 'a=b=c'\n"
141
+ "print('ok')"
142
+ )
143
+ mock_runtime.execute_code.assert_any_call(expected, output_hook=ANY, timeout=30.0)
144
+
145
+
146
+def test_cli_exec_env_escapes_tricky_literals(
147
+ mock_store, mock_runtime_class, mock_common_state
148
+):
149
+ mock_session = MagicMock()
150
+ mock_session.name = "s1"
151
+ mock_session.url = "http://url"
152
+ mock_session.token = "token"
153
+ mock_session.kernel_id = None
154
+ mock_session.session_id = None
155
+ mock_store.get.return_value = mock_session
156
+
157
+ mock_common_state.resolve_session.return_value = "s1"
158
+ mock_runtime = mock_runtime_class.return_value
159
+ mock_runtime.execute_code.return_value = []
160
+
161
+ value = "quote'back\\slash=µ"
162
+ result = runner.invoke(
163
+ app, ["exec", "-s", "s1", "--env", f"TRICKY={value}"], input="print('ok')"
164
+ )
165
+
166
+ assert result.exit_code == 0, result.output
167
+ expected = f"import os\nos.environ['TRICKY'] = {value!r}\nprint('ok')"
168
+ mock_runtime.execute_code.assert_any_call(expected, output_hook=ANY, timeout=30.0)
169
+
170
+
171
+def test_cli_exec_malformed_env_errors_before_session_resolution(
172
+ mock_runtime_class, mock_common_state
173
+):
174
+ result = runner.invoke(
175
+ app, ["exec", "-s", "s1", "--env", "HF_TOKEN"], input="print('ok')"
176
+ )
177
+
178
+ assert result.exit_code != 0
179
+ assert "Expected KEY=VALUE" in result.output
180
+ mock_common_state.resolve_session.assert_not_called()
181
+ mock_runtime_class.assert_not_called()
182
+
183
+
184
def test_cli_exec_not_found(mock_common_state):
185
# Case where resolve_session fails
186
mock_common_state.resolve_session.side_effect = SystemExit(1)
tests/test_run.py
+108
@@ -186,6 +186,105 @@ def test_run_passes_argv(
186
assert "'--flag-for-script'" in body
187
188
189
+def test_run_env_flag_after_script_sets_env_and_preserves_argv(
190
+ mock_client,
191
+ mock_store,
192
+ mock_runtime_class,
193
+ mock_spawn_keep_alive,
194
+ assign_response,
195
+ script_path,
196
+):
197
+ """`--env KEY=VALUE` after the script path should configure the remote
198
+ environment, not get forwarded into sys.argv."""
199
+ mock_client.assign.return_value = assign_response
200
+ mock_runtime = mock_runtime_class.return_value
201
+ mock_runtime.execute_code.return_value = []
202
+
203
+ persisted = {}
204
+ mock_store.add.side_effect = lambda s: persisted.setdefault("s", s)
205
+ mock_store.get.side_effect = lambda name: persisted.get("s")
206
+
207
+ result = runner.invoke(
208
+ app, ["run", str(script_path), "--env", "HF_TOKEN=abc", "alpha"]
209
+ )
210
+
211
+ assert result.exit_code == 0, result.output
212
+ code_calls = [c.args[0] for c in mock_runtime.execute_code.call_args_list]
213
+ body = next(c for c in code_calls if "hello from script" in c)
214
+ assert "import os\n" in body
215
+ assert "os.environ['HF_TOKEN'] = 'abc'" in body
216
+ assert body.index("os.environ['HF_TOKEN'] = 'abc'") < body.index(
217
+ "print('hello from script')"
218
+ )
219
+ assert "'alpha'" in body
220
+ assert "'--env'" not in body
221
+ assert "'HF_TOKEN=abc'" not in body
222
+
223
+
224
+def test_run_env_flags_accumulate_and_split_on_first_equals(
225
+ mock_client,
226
+ mock_store,
227
+ mock_runtime_class,
228
+ mock_spawn_keep_alive,
229
+ assign_response,
230
+ script_path,
231
+):
232
+ """Repeated env flags should all become assignments; values may contain
233
+ additional '=' characters."""
234
+ mock_client.assign.return_value = assign_response
235
+ mock_runtime = mock_runtime_class.return_value
236
+ mock_runtime.execute_code.return_value = []
237
+
238
+ persisted = {}
239
+ mock_store.add.side_effect = lambda s: persisted.setdefault("s", s)
240
+ mock_store.get.side_effect = lambda name: persisted.get("s")
241
+
242
+ result = runner.invoke(
243
+ app,
244
+ [
245
+ "run",
246
+ "--env",
247
+ "HF_TOKEN=abc",
248
+ "--env",
249
+ "B64=a=b=c",
250
+ str(script_path),
251
+ ],
252
+ )
253
+
254
+ assert result.exit_code == 0, result.output
255
+ code_calls = [c.args[0] for c in mock_runtime.execute_code.call_args_list]
256
+ body = next(c for c in code_calls if "hello from script" in c)
257
+ assert "os.environ['HF_TOKEN'] = 'abc'" in body
258
+ assert "os.environ['B64'] = 'a=b=c'" in body
259
+
260
+
261
+def test_run_env_flag_escapes_tricky_literals(
262
+ mock_client,
263
+ mock_store,
264
+ mock_runtime_class,
265
+ mock_spawn_keep_alive,
266
+ assign_response,
267
+ script_path,
268
+):
269
+ """Quotes, backslashes, '=' and non-ASCII values should round-trip as safe
270
+ Python literals."""
271
+ mock_client.assign.return_value = assign_response
272
+ mock_runtime = mock_runtime_class.return_value
273
+ mock_runtime.execute_code.return_value = []
274
+
275
+ persisted = {}
276
+ mock_store.add.side_effect = lambda s: persisted.setdefault("s", s)
277
+ mock_store.get.side_effect = lambda name: persisted.get("s")
278
+
279
+ value = "quote'back\\slash=µ"
280
+ result = runner.invoke(app, ["run", "--env", f"TRICKY={value}", str(script_path)])
281
+
282
+ assert result.exit_code == 0, result.output
283
+ code_calls = [c.args[0] for c in mock_runtime.execute_code.call_args_list]
284
+ body = next(c for c in code_calls if "hello from script" in c)
285
+ assert f"os.environ['TRICKY'] = {value!r}" in body
286
+
287
+
288
def test_run_sets_dunder_main(
289
mock_client,
290
mock_store,
@@ -352,6 +451,15 @@ def test_run_nonexistent_script_errors_before_assign(mock_client):
451
mock_client.assign.assert_not_called()
452
453
454
+def test_run_malformed_env_errors_before_assign(mock_client, script_path):
455
+ """Malformed env entries must be rejected locally before any VM allocation."""
456
+ result = runner.invoke(app, ["run", str(script_path), "--env", "HF_TOKEN"])
457
+
458
+ assert result.exit_code != 0
459
+ assert "Expected KEY=VALUE" in result.output
460
+ mock_client.assign.assert_not_called()
461
+
462
+
463
# ---------------------------------------------------------------------------
464
# SystemExit handling — the kernel reports `sys.exit(N)` as an error output of
465
# `ename=='SystemExit'`. We want native-`python`-like semantics: exit 0 for