Skip to content

Commit 42361d2

Browse files
Simplify exit code test command
1 parent 677cfef commit 42361d2

2 files changed

Lines changed: 431 additions & 0 deletions

File tree

scripts/virtual_container.py

Lines changed: 273 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,273 @@
1+
#!/usr/bin/env python3
2+
"""Utility to spin up a lightweight virtual container for running commands.
3+
4+
The virtual container concept implemented here provides an isolated working
5+
directory (with environment settings) for executing shell commands. It is
6+
aimed at cases where a throwaway sandbox is preferred to running commands
7+
directly in the repository tree. The tool supports both interactive and
8+
one-off command execution modes.
9+
"""
10+
11+
from __future__ import annotations
12+
13+
import argparse
14+
import os
15+
import shutil
16+
import subprocess
17+
import sys
18+
import tempfile
19+
from dataclasses import dataclass
20+
from pathlib import Path
21+
from typing import Dict, Iterable, Mapping, MutableMapping, Sequence
22+
23+
24+
@dataclass
25+
class VirtualContainer:
26+
"""Manage a throwaway workspace for executing shell commands.
27+
28+
Parameters
29+
----------
30+
root:
31+
The root directory that represents ``/`` within the virtual container.
32+
cwd:
33+
The current working directory inside the container. Defaults to the
34+
``root``.
35+
env:
36+
A copy of the environment variables that should be supplied to spawned
37+
commands.
38+
"""
39+
40+
root: Path
41+
cwd: Path
42+
env: MutableMapping[str, str]
43+
44+
@classmethod
45+
def create(
46+
cls,
47+
*,
48+
base_dir: str | None = None,
49+
env: Mapping[str, str] | None = None,
50+
) -> "VirtualContainer":
51+
"""Instantiate a new virtual container."""
52+
53+
if base_dir is not None:
54+
Path(base_dir).mkdir(parents=True, exist_ok=True)
55+
root_dir = Path(
56+
tempfile.mkdtemp(prefix="virtual_container_", dir=base_dir)
57+
).resolve()
58+
env_vars: Dict[str, str] = dict(os.environ)
59+
if env:
60+
env_vars.update(env)
61+
return cls(root=root_dir, cwd=root_dir, env=env_vars)
62+
63+
def run(
64+
self,
65+
command: Sequence[str] | str,
66+
*,
67+
check: bool = False,
68+
) -> subprocess.CompletedProcess[str]:
69+
"""Run a command inside the container and return the completed process."""
70+
71+
if isinstance(command, str):
72+
exec_command = command
73+
shell = True
74+
else:
75+
exec_command = list(command)
76+
shell = False
77+
completed = subprocess.run(
78+
exec_command,
79+
cwd=str(self.cwd),
80+
env=self.env,
81+
capture_output=True,
82+
text=True,
83+
shell=shell,
84+
)
85+
if check and completed.returncode != 0:
86+
raise subprocess.CalledProcessError(
87+
completed.returncode,
88+
exec_command,
89+
output=completed.stdout,
90+
stderr=completed.stderr,
91+
)
92+
return completed
93+
94+
def cleanup(self) -> None:
95+
"""Remove the virtual container directory tree."""
96+
97+
if self.root.exists():
98+
shutil.rmtree(self.root)
99+
100+
def format_path(self, path: Path | None = None) -> str:
101+
"""Return a container-relative representation of *path*."""
102+
103+
target = path or self.cwd
104+
if target == self.root:
105+
return "/"
106+
return f"/{target.relative_to(self.root)}"
107+
108+
def resolve_path(self, raw_path: str) -> Path:
109+
"""Resolve *raw_path* to a directory within the container.
110+
111+
Absolute paths are interpreted relative to the container root in the
112+
same way they would be inside a real container.
113+
"""
114+
115+
if not raw_path:
116+
raise ValueError("path cannot be empty")
117+
path_obj = Path(raw_path)
118+
if path_obj.is_absolute():
119+
normalized = self.root.joinpath(*path_obj.parts[1:])
120+
else:
121+
normalized = self.cwd / path_obj
122+
resolved = normalized.resolve()
123+
try:
124+
resolved.relative_to(self.root)
125+
except ValueError as exc: # pragma: no cover - safety guard
126+
raise ValueError("cannot leave the container root") from exc
127+
if not resolved.exists():
128+
raise FileNotFoundError(raw_path)
129+
if not resolved.is_dir():
130+
raise NotADirectoryError(raw_path)
131+
return resolved
132+
133+
def change_directory(self, raw_path: str) -> Path:
134+
"""Change the container's working directory."""
135+
136+
target = self.resolve_path(raw_path)
137+
self.cwd = target
138+
return self.cwd
139+
140+
141+
def parse_env_assignments(assignments: Iterable[str]) -> Dict[str, str]:
142+
"""Parse ``KEY=VALUE`` assignments supplied on the command line."""
143+
144+
parsed: Dict[str, str] = {}
145+
for item in assignments:
146+
if "=" not in item:
147+
raise ValueError(f"invalid environment assignment: '{item}'")
148+
key, value = item.split("=", 1)
149+
key = key.strip()
150+
if not key:
151+
raise ValueError("environment variable name cannot be empty")
152+
parsed[key] = value
153+
return parsed
154+
155+
156+
def build_parser() -> argparse.ArgumentParser:
157+
"""Create the command-line argument parser for the script."""
158+
159+
parser = argparse.ArgumentParser(
160+
description=(
161+
"Create a lightweight virtual container directory and execute commands "
162+
"inside it."
163+
)
164+
)
165+
parser.add_argument(
166+
"command",
167+
nargs=argparse.REMAINDER,
168+
help=(
169+
"Command to execute inside the container. When omitted an interactive "
170+
"prompt is started."
171+
),
172+
)
173+
parser.add_argument(
174+
"-e",
175+
"--env",
176+
action="append",
177+
default=[],
178+
metavar="KEY=VALUE",
179+
help="Environment variables to inject into the container.",
180+
)
181+
parser.add_argument(
182+
"--base-dir",
183+
type=str,
184+
default=None,
185+
help="Directory where container instances should be created.",
186+
)
187+
parser.add_argument(
188+
"--persist",
189+
action="store_true",
190+
help="Do not delete the container directory when the program exits.",
191+
)
192+
return parser
193+
194+
195+
def container_prompt(container: VirtualContainer) -> str:
196+
"""Return a shell-like prompt for the interactive session."""
197+
198+
return f"{container.format_path()}$ "
199+
200+
201+
def interactive_loop(container: VirtualContainer) -> int:
202+
"""Run an interactive terminal loop inside the container."""
203+
204+
print("Starting virtual container interactive session.")
205+
print("Type 'exit' or 'quit' (or press Ctrl-D) to leave.")
206+
while True:
207+
try:
208+
command = input(container_prompt(container)).strip()
209+
except EOFError:
210+
print()
211+
break
212+
if not command:
213+
continue
214+
if command in {"exit", "quit"}:
215+
break
216+
if command == "cd":
217+
container.change_directory("/")
218+
continue
219+
if command.startswith("cd "):
220+
target = command[3:].strip() or "/"
221+
try:
222+
container.change_directory(target)
223+
except OSError as exc:
224+
print(f"cd: {exc}")
225+
continue
226+
result = container.run(command)
227+
if result.stdout:
228+
sys.stdout.write(result.stdout)
229+
if result.stderr:
230+
sys.stderr.write(result.stderr)
231+
print(f"[exit {result.returncode}]")
232+
return 0
233+
234+
235+
def main(argv: Sequence[str] | None = None) -> int:
236+
parser = build_parser()
237+
args = parser.parse_args(argv)
238+
try:
239+
env_vars = parse_env_assignments(args.env)
240+
except ValueError as exc:
241+
parser.error(str(exc))
242+
container = VirtualContainer.create(base_dir=args.base_dir, env=env_vars)
243+
exit_code = 0
244+
try:
245+
if args.command:
246+
command = args.command
247+
if command and command[0] == "--":
248+
command = command[1:]
249+
if not command:
250+
parser.error("a command must follow '--'")
251+
if len(command) == 1:
252+
command = command[0]
253+
result = container.run(command)
254+
if result.stdout:
255+
sys.stdout.write(result.stdout)
256+
if result.stderr:
257+
sys.stderr.write(result.stderr)
258+
exit_code = result.returncode
259+
else:
260+
exit_code = interactive_loop(container)
261+
finally:
262+
if args.persist:
263+
print(
264+
"Container preserved at",
265+
container.root,
266+
)
267+
else:
268+
container.cleanup()
269+
return exit_code
270+
271+
272+
if __name__ == "__main__":
273+
raise SystemExit(main())

0 commit comments

Comments
 (0)