locki 0.0.1__tar.gz

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.
locki-0.0.1/PKG-INFO ADDED
@@ -0,0 +1,12 @@
1
+ Metadata-Version: 2.3
2
+ Name: locki
3
+ Version: 0.0.1
4
+ Summary: Lock down agents in a VM, enabling mischief without consequences
5
+ Requires-Dist: anyio>=4.12.1
6
+ Requires-Dist: pydantic>=2.12.5
7
+ Requires-Dist: rich>=14.3.3
8
+ Requires-Dist: typer>=0.24.1
9
+ Requires-Python: >=3.14, <3.15
10
+ Description-Content-Type: text/markdown
11
+
12
+ # locki
locki-0.0.1/README.md ADDED
@@ -0,0 +1 @@
1
+ # locki
@@ -0,0 +1,34 @@
1
+ [project]
2
+ name = "locki"
3
+ version = "0.0.1"
4
+ description = "Lock down agents in a VM, enabling mischief without consequences"
5
+ readme = "README.md"
6
+ requires-python = ">=3.14,<3.15"
7
+ dependencies = [
8
+ "anyio>=4.12.1",
9
+ "pydantic>=2.12.5",
10
+ "rich>=14.3.3",
11
+ "typer>=0.24.1",
12
+ ]
13
+
14
+ [dependency-groups]
15
+ dev = [
16
+ "ruff>=0.15.7",
17
+ "wheel>=0.46.3",
18
+ ]
19
+
20
+ [project.scripts]
21
+ locki = "locki:app"
22
+
23
+ [build-system]
24
+ requires = ["uv_build>=0.10.0,<0.11.0"]
25
+ build-backend = "uv_build"
26
+
27
+ [tool.ruff]
28
+ line-length = 120
29
+ target-version = "py314"
30
+ lint.select = [
31
+ "E", "W", "F", "UP", "I", "B", "N", "C4", "Q", "SIM", "RUF", "TID", "ASYNC",
32
+ ]
33
+ lint.ignore = ["E501"]
34
+ force-exclude = true
@@ -0,0 +1,350 @@
1
+ import functools
2
+ import importlib.resources
3
+ import os
4
+ import pathlib
5
+ import secrets
6
+ import shlex
7
+ import shutil
8
+ import subprocess
9
+ import sys
10
+ import typing
11
+
12
+ import typer
13
+
14
+ from locki.async_typer import AsyncTyperWithAliases
15
+ from locki.config import load_config
16
+ from locki.console import console
17
+ from locki.utils import run_command, verbosity
18
+
19
+ app = AsyncTyperWithAliases(
20
+ name="locki",
21
+ help="Lima VM wrapper that protects worktrees by offering isolated execution environments.",
22
+ no_args_is_help=True,
23
+ )
24
+
25
+ LOCKI_HOME = pathlib.Path.home() / ".locki"
26
+ LIMA_HOME = LOCKI_HOME / "lima"
27
+ WORKTREES_HOME = LOCKI_HOME / "worktrees"
28
+
29
+
30
+ @functools.cache
31
+ def limactl() -> str:
32
+ bundled = importlib.resources.files("locki") / "data" / "bin" / "limactl"
33
+ if bundled.is_file():
34
+ return str(bundled)
35
+ system = shutil.which("limactl")
36
+ if system:
37
+ return system
38
+ console.error("limactl is not installed. Please install Lima or use a platform-specific locki wheel.")
39
+ sys.exit(1)
40
+
41
+
42
+ async def run_in_vm(
43
+ command: list[str],
44
+ message: str,
45
+ env: dict[str, str] | None = None,
46
+ input: bytes | None = None,
47
+ check: bool = True,
48
+ ) -> subprocess.CompletedProcess[bytes]:
49
+ return await run_command(
50
+ [limactl(), "shell", "--start", "--preserve-env", "--tty=false", "locki", "--", "sudo", "-E", *command],
51
+ message,
52
+ env={"LIMA_HOME": str(LIMA_HOME)} | (env or {}),
53
+ cwd="/",
54
+ input=input,
55
+ check=check,
56
+ )
57
+
58
+
59
+ async def ensure_vm() -> None:
60
+ LOCKI_HOME.mkdir(exist_ok=True)
61
+ LIMA_HOME.mkdir(exist_ok=True, parents=True)
62
+ WORKTREES_HOME.mkdir(parents=True, exist_ok=True)
63
+ await run_command(
64
+ [
65
+ limactl(),
66
+ "--tty=false",
67
+ "create",
68
+ str(importlib.resources.files("locki").joinpath("locki.yaml")),
69
+ "--mount-writable",
70
+ "--name=locki",
71
+ ],
72
+ "Preparing VM",
73
+ env={"LIMA_HOME": str(LIMA_HOME)},
74
+ cwd="/",
75
+ check=False,
76
+ )
77
+ await run_command(
78
+ [
79
+ limactl(),
80
+ "--tty=false",
81
+ "start",
82
+ "locki",
83
+ ],
84
+ "Starting VM",
85
+ env={"LIMA_HOME": str(LIMA_HOME)},
86
+ cwd="/",
87
+ check=False,
88
+ )
89
+
90
+
91
+ @functools.cache
92
+ def git_root() -> pathlib.Path:
93
+ current = pathlib.Path.cwd()
94
+ while True:
95
+ dot_git = current / ".git"
96
+ if dot_git.is_dir():
97
+ return current
98
+ if dot_git.is_file():
99
+ content = dot_git.read_text().strip()
100
+ if content.startswith("gitdir:"):
101
+ wt_gitdir = pathlib.Path(content.split(":", 1)[1].strip())
102
+ if not wt_gitdir.is_absolute():
103
+ wt_gitdir = (current / wt_gitdir).resolve()
104
+ main_git_dir = (wt_gitdir / ".." / "..").resolve()
105
+ if main_git_dir.name == ".git":
106
+ return main_git_dir.parent
107
+ return current
108
+ if current.parent == current:
109
+ console.error("Not inside a git repository.")
110
+ sys.exit(1)
111
+ current = current.parent
112
+
113
+
114
+ async def find_worktree_for_branch(branch: str) -> pathlib.Path | None:
115
+ """Return the worktree path for a branch managed by locki, or None."""
116
+ result = await run_command(
117
+ ["git", "-C", str(git_root()), "worktree", "list", "--porcelain"],
118
+ "Listing worktrees",
119
+ )
120
+ current_path: pathlib.Path | None = None
121
+ for line in result.stdout.decode().splitlines():
122
+ if line.startswith("worktree "):
123
+ current_path = pathlib.Path(line.split(" ", 1)[1])
124
+ elif (
125
+ line.startswith("branch refs/heads/")
126
+ and line.split("/")[-1] == branch
127
+ and current_path
128
+ and current_path.is_relative_to(WORKTREES_HOME)
129
+ ):
130
+ return current_path
131
+ return None
132
+
133
+
134
+ async def ensure_worktree(branch: str) -> pathlib.Path:
135
+ """Ensure a locki-managed worktree exists for the branch. Returns the worktree path."""
136
+ existing = await find_worktree_for_branch(branch)
137
+ if existing:
138
+ return existing
139
+
140
+ await run_command(
141
+ ["git", "-C", str(git_root()), "worktree", "prune"],
142
+ "Pruning stale git worktrees",
143
+ )
144
+
145
+ repo_name = git_root().name.replace("/", "-").replace(".", "-").lower()
146
+ safe_branch = branch.replace("/", "-").replace(".", "-").lower()
147
+ wt_id = f"{repo_name}--{safe_branch}--{secrets.token_hex(4)}"
148
+ wt_path = WORKTREES_HOME / wt_id
149
+ wt_path.mkdir(parents=True, exist_ok=True)
150
+
151
+ result = await run_command(
152
+ ["git", "-C", str(git_root()), "rev-parse", "--verify", f"refs/heads/{branch}"],
153
+ f"Checking if branch '{branch}' exists",
154
+ check=False,
155
+ )
156
+ if result.returncode != 0:
157
+ await run_command(
158
+ ["git", "-C", str(git_root()), "branch", branch],
159
+ f"Creating branch '{branch}'",
160
+ )
161
+
162
+ await run_command(
163
+ ["git", "-C", str(git_root()), "worktree", "add", str(wt_path), branch],
164
+ f"Creating worktree for '{branch}'",
165
+ )
166
+
167
+ return wt_path
168
+
169
+
170
+ async def ensure_container(wt_id: str, wt_path: pathlib.Path, config) -> None:
171
+ """Ensure an Incus container exists for the given worktree (idempotent)."""
172
+ result = await run_in_vm(
173
+ ["incus", "list", "--format=csv", "--columns=n", wt_id],
174
+ "Checking container",
175
+ check=False,
176
+ )
177
+ if wt_id in result.stdout.decode():
178
+ await run_in_vm(
179
+ ["incus", "start", wt_id],
180
+ "Starting container",
181
+ check=False,
182
+ )
183
+ return
184
+
185
+ incus_image = config.get_incus_image()
186
+
187
+ local_path = git_root() / incus_image
188
+ if local_path.is_file():
189
+ local_file = local_path.resolve()
190
+ await run_command(
191
+ [limactl(), "copy", str(local_file), "locki:/tmp/image"],
192
+ "Copying image into VM",
193
+ env={"LIMA_HOME": str(LIMA_HOME)},
194
+ cwd="/",
195
+ )
196
+ await run_in_vm(
197
+ ["incus", "image", "import", "/tmp/image", f"--alias={wt_id}"],
198
+ "Importing container image",
199
+ )
200
+ await run_in_vm(["rm", "-f", "/tmp/image"], "Cleaning up image file", check=False)
201
+ image_ref = wt_id
202
+ else:
203
+ image_ref = incus_image
204
+
205
+ await run_in_vm(
206
+ ["incus", "init", image_ref, wt_id],
207
+ "Creating container",
208
+ )
209
+
210
+ if local_path.is_file():
211
+ await run_in_vm(
212
+ ["incus", "image", "delete", wt_id],
213
+ "Cleaning up imported image",
214
+ check=False,
215
+ )
216
+
217
+ await run_in_vm(
218
+ [
219
+ "incus",
220
+ "config",
221
+ "device",
222
+ "add",
223
+ wt_id,
224
+ "worktree",
225
+ "disk",
226
+ f"source={wt_path}",
227
+ f"path={wt_path}",
228
+ ],
229
+ "Mounting worktree into container",
230
+ )
231
+
232
+ await run_in_vm(
233
+ ["incus", "start", wt_id],
234
+ "Starting container",
235
+ )
236
+
237
+
238
+ @app.command("shell", help="Open a shell in the per-branch container (creates branch/worktree/container if needed).")
239
+ async def shell_cmd(
240
+ branch: typing.Annotated[str, typer.Argument(help="Branch name to work on")],
241
+ verbose: typing.Annotated[bool, typer.Option("-v", "--verbose", help="Show verbose output")] = False,
242
+ ):
243
+ with verbosity(verbose):
244
+ git_root() # fail fast if not in a git repo
245
+
246
+ await ensure_vm()
247
+
248
+ wt_path = await ensure_worktree(branch)
249
+ wt_id = wt_path.relative_to(WORKTREES_HOME).parts[0]
250
+
251
+ config = load_config(git_root())
252
+ await ensure_container(wt_id, wt_path, config)
253
+
254
+ forwarded_env = {"TERM", "COLORTERM", "TERM_PROGRAM", "TERM_PROGRAM_VERSION", "LANG", "SSH_TTY"}
255
+
256
+ os.environ["LIMA_HOME"] = str(LIMA_HOME)
257
+ os.environ["LIMA_SHELLENV_ALLOW"] = ",".join(forwarded_env)
258
+
259
+ os.execvp(
260
+ limactl(),
261
+ [
262
+ limactl(),
263
+ "shell",
264
+ "--yes",
265
+ "--preserve-env",
266
+ "--start",
267
+ "locki",
268
+ "--",
269
+ "bash",
270
+ "-c",
271
+ " ".join(
272
+ [
273
+ "sudo",
274
+ "incus",
275
+ "exec",
276
+ shlex.quote(wt_id),
277
+ "--cwd",
278
+ shlex.quote(str(wt_path)),
279
+ *(f"--env={env}=${env}" for env in forwarded_env),
280
+ "--",
281
+ "bash",
282
+ "--login",
283
+ ]
284
+ ),
285
+ ],
286
+ )
287
+
288
+
289
+ @app.command("remove", help="Remove a branch's worktree and container.")
290
+ async def remove_cmd(
291
+ branch: typing.Annotated[str, typer.Argument(help="Branch name to remove")],
292
+ force: typing.Annotated[bool, typer.Option("--force", "-f", help="Skip safety checks")] = False,
293
+ verbose: typing.Annotated[bool, typer.Option("-v", "--verbose", help="Show verbose output")] = False,
294
+ ):
295
+ with verbosity(verbose):
296
+ wt_path = await find_worktree_for_branch(branch)
297
+
298
+ if wt_path is None:
299
+ console.info(f"No locki-managed worktree found for '{branch}', nothing to do.")
300
+ return
301
+
302
+ if not force and (await run_command(
303
+ ["git", "-C", str(wt_path), "status", "--porcelain"],
304
+ "Checking for uncommitted changes",
305
+ check=False,
306
+ )).stdout.strip():
307
+ console.error(f"Worktree for {branch} in {wt_path} has uncommitted changes. Commit or stash them, or use --force.")
308
+ sys.exit(1)
309
+
310
+ wt_id = wt_path.relative_to(WORKTREES_HOME).parts[0]
311
+
312
+ await run_in_vm(
313
+ ["incus", "delete", "--force", wt_id],
314
+ "Deleting container",
315
+ check=False,
316
+ )
317
+
318
+ await run_command(
319
+ ["git", "-C", str(git_root()), "worktree", "remove", "--force", str(wt_path)],
320
+ "Removing worktree",
321
+ check=False,
322
+ )
323
+
324
+
325
+ @app.command("list", help="List branches with locki-managed worktrees.")
326
+ async def list_cmd(
327
+ verbose: typing.Annotated[bool, typer.Option("-v", "--verbose", help="Show verbose output")] = False,
328
+ ):
329
+ with verbosity(verbose, show_success_status=False):
330
+ result = await run_command(
331
+ ["git", "-C", str(git_root()), "worktree", "list", "--porcelain"],
332
+ "Listing worktrees",
333
+ )
334
+
335
+ found = False
336
+ current_path: pathlib.Path | None = None
337
+ current_branch: str | None = None
338
+ for line in result.stdout.decode().splitlines():
339
+ if line.startswith("worktree "):
340
+ current_path = pathlib.Path(line.split(" ", 1)[1])
341
+ current_branch = None
342
+ elif line.startswith("branch refs/heads/"):
343
+ current_branch = line.removeprefix("branch refs/heads/")
344
+ elif line == "" and current_path and current_branch:
345
+ if current_path.is_relative_to(WORKTREES_HOME):
346
+ console.print(f"{current_branch} [dim]{current_path}[/dim]")
347
+ found = True
348
+
349
+ if not found:
350
+ console.info("No locki-managed worktrees found.")
@@ -0,0 +1,56 @@
1
+ import asyncio
2
+ import functools
3
+ import inspect
4
+ import sys
5
+
6
+ import typer
7
+ from typer.core import TyperGroup
8
+
9
+ from locki.console import err_console
10
+
11
+
12
+ class AsyncTyper(typer.Typer):
13
+ def command(self, *args, **kwargs):
14
+ parent_decorator = super().command(*args, **kwargs)
15
+
16
+ def decorator(f):
17
+ @functools.wraps(f)
18
+ def wrapped_f(*args, **kwargs):
19
+ if sys.stdout.isatty():
20
+ sys.stdout.write("\x1b[>0u")
21
+ sys.stdout.flush()
22
+ try:
23
+ if inspect.iscoroutinefunction(f):
24
+ return asyncio.run(f(*args, **kwargs))
25
+ else:
26
+ return f(*args, **kwargs)
27
+ except* Exception as eg:
28
+ for exc in eg.exceptions:
29
+ err_console.error(f"{type(exc).__name__}: {exc}")
30
+ sys.exit(1)
31
+ finally:
32
+ if sys.stdout.isatty():
33
+ sys.stdout.write("\x1b[<u")
34
+ sys.stdout.flush()
35
+
36
+ parent_decorator(wrapped_f)
37
+ return f
38
+
39
+ return decorator
40
+
41
+
42
+ class AliasGroup(TyperGroup):
43
+ """Support comma/pipe-separated command name aliases, e.g. 'start|up'."""
44
+
45
+ def get_command(self, ctx, cmd_name):
46
+ for cmd in self.commands.values():
47
+ if cmd.name and cmd_name in cmd.name.replace(" ", "").split(","):
48
+ cmd_name = cmd.name
49
+ break
50
+ return super().get_command(ctx, cmd_name)
51
+
52
+
53
+ class AsyncTyperWithAliases(AsyncTyper):
54
+ def __init__(self, *args, **kwargs):
55
+ kwargs.setdefault("cls", AliasGroup)
56
+ super().__init__(*args, **kwargs)
@@ -0,0 +1,38 @@
1
+ import pathlib
2
+ import platform
3
+ import sys
4
+ import tomllib
5
+
6
+ import pydantic
7
+
8
+ from locki.console import console
9
+
10
+ DEFAULT_INCUS_IMAGES: dict[str, str] = {
11
+ "arm64": "locki-base",
12
+ "amd64": "locki-base",
13
+ }
14
+
15
+
16
+ class LockiConfig(pydantic.BaseModel):
17
+ incus_image: dict[str, str] = pydantic.Field(default_factory=lambda: dict(DEFAULT_INCUS_IMAGES))
18
+
19
+ def get_incus_image(self) -> str:
20
+ arch = platform.machine()
21
+ if arch not in self.incus_image:
22
+ console.error(
23
+ f"No incus_image configured for architecture '{arch}'. Available: {', '.join(self.incus_image)}"
24
+ )
25
+ sys.exit(1)
26
+ return self.incus_image[arch]
27
+
28
+
29
+ def load_config(git_root: pathlib.Path) -> LockiConfig:
30
+ config_path = git_root / "locki.toml"
31
+ if not config_path.exists():
32
+ return LockiConfig()
33
+ try:
34
+ with open(config_path, "rb") as f:
35
+ return LockiConfig.model_validate(tomllib.load(f))
36
+ except (tomllib.TOMLDecodeError, pydantic.ValidationError) as e:
37
+ console.error(f"Invalid locki.toml: {e}")
38
+ sys.exit(1)
@@ -0,0 +1,22 @@
1
+ from rich.console import Console
2
+
3
+
4
+ class ExtendedConsole(Console):
5
+ def error(self, message: str):
6
+ self.print(f":boom: [bold red]ERROR[/bold red]: {message}")
7
+
8
+ def warning(self, message: str):
9
+ self.print(f":warning: [yellow]WARNING[/yellow]: {message}")
10
+
11
+ def hint(self, message: str):
12
+ self.print(f":bulb: [bright_cyan]HINT[/bright_cyan]: {message}")
13
+
14
+ def success(self, message: str):
15
+ self.print(f":white_check_mark: [green]SUCCESS[/green]: {message}")
16
+
17
+ def info(self, message: str):
18
+ self.print(f":memo: INFO: {message}")
19
+
20
+
21
+ err_console = ExtendedConsole(stderr=True)
22
+ console = ExtendedConsole()
@@ -0,0 +1,115 @@
1
+ minimumLimaVersion: "2.0.0"
2
+
3
+ base:
4
+ - template:fedora
5
+
6
+ containerd:
7
+ system: false
8
+ user: false
9
+
10
+ mounts:
11
+ - location: "~/.locki/worktrees"
12
+ writable: true
13
+
14
+ provision:
15
+ - mode: system
16
+ script: |
17
+ #!/bin/bash
18
+ set -euxo pipefail
19
+ if command -v incus; then exit 0; fi
20
+ echo "root:1000000:1000000000" >> /etc/subuid
21
+ echo "root:1000000:1000000000" >> /etc/subgid
22
+ dnf install -y --setopt install_weak_deps=False incus incus-client
23
+ systemctl enable --now incus
24
+ mkdir -p /var/cache/locki
25
+ incus admin init --preseed <<EOF
26
+ storage_pools:
27
+ - name: default
28
+ driver: dir
29
+ networks:
30
+ - name: incusbr0
31
+ type: bridge
32
+ config:
33
+ ipv4.address: 10.99.0.1/24
34
+ ipv4.nat: "true"
35
+ ipv6.address: none
36
+ profiles:
37
+ - name: default
38
+ config:
39
+ environment.PATH: /root/.local/bin:/opt/mise/shims:/var/cache/locki/pnpm:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
40
+ environment.PNPM_HOME: /var/cache/locki/pnpm
41
+ environment.UV_CACHE_DIR: /var/cache/locki/uv
42
+ environment.COREPACK_ENABLE_DOWNLOAD_PROMPT: 0
43
+ environment.MISE_DATA_DIR: /opt/mise
44
+ environment.MISE_CONFIG_DIR: /opt/mise
45
+ environment.MISE_CACHE_DIR: /var/cache/locki/mise
46
+ environment.MISE_INSTALL_PATH: /usr/local/bin/mise
47
+ environment.MISE_TRUSTED_CONFIG_PATHS: /
48
+ environment.IS_SANDBOX: 1
49
+ security.nesting: "true"
50
+ security.privileged: "true"
51
+ raw.lxc: |
52
+ lxc.mount.auto = proc:rw sys:rw
53
+ lxc.cap.drop =
54
+ devices:
55
+ root:
56
+ path: /
57
+ pool: default
58
+ type: disk
59
+ eth0:
60
+ name: eth0
61
+ network: incusbr0
62
+ type: nic
63
+ kmsg:
64
+ path: /dev/kmsg
65
+ source: /dev/kmsg
66
+ type: unix-char
67
+ cache:
68
+ path: /var/cache/locki
69
+ source: /var/cache/locki
70
+ type: disk
71
+ claude-user:
72
+ path: /root/.claude
73
+ source: /root/.claude
74
+ type: disk
75
+ claude-managed:
76
+ path: /etc/claude-code
77
+ source: /etc/claude-code
78
+ type: disk
79
+ EOF
80
+ - mode: system
81
+ script: |
82
+ #!/bin/bash
83
+ set -euxo pipefail
84
+ if incus image list --format csv | grep -q locki-base; then exit 0; fi
85
+ mkdir -p /root/.claude
86
+ incus launch images:fedora/43 locki-base-builder
87
+ incus exec locki-base-builder -- bash <<'ENDBUILD'
88
+ set -euxo pipefail
89
+ curl --retry 10 --retry-all-errors https://mirrors.fedoraproject.org
90
+ dnf install -y \
91
+ https://mirrors.rpmfusion.org/free/fedora/rpmfusion-free-release-$(rpm -E %fedora).noarch.rpm \
92
+ https://mirrors.rpmfusion.org/nonfree/fedora/rpmfusion-nonfree-release-$(rpm -E %fedora).noarch.rpm
93
+ curl https://mise.run | sh
94
+ mise use -g claude@latest
95
+ ENDBUILD
96
+ incus stop locki-base-builder
97
+ incus publish locki-base-builder --alias locki-base
98
+ incus delete locki-base-builder
99
+ - mode: data
100
+ path: /etc/claude-code/managed-settings.json
101
+ content: |
102
+ {
103
+ "skipDangerousModePermissionPrompt": true,
104
+ "permissions": {
105
+ "defaultMode": "bypassPermissions"
106
+ }
107
+ }
108
+ - mode: data
109
+ path: /etc/claude-code/CLAUDE.md
110
+ content: |
111
+ You are running inside a VM, in a managed Incus container. This is an ephemeral environment designed to keep the main machine safe from malfunctioning agents. Before delving into your task, start by setting up the environment:
112
+
113
+ Container has Mise-en-place preinstalled. If there is `mise.toml` in the repo root, run `mise install` to install tools from it. Otherwise, check files like `.tool-versions`, `.nvmrc`, `pyproject.toml`, etc. or docs like `README.md`, `CONTRIBUTING.md` to determine needed tools and versions. Enable specific tool versions using e.g.: `mise use -g python@3.12.1`, `mise use -g node@22`, `mise use -g jq`. Container has Fedora 43 with RPM Fusion enabled -- install packages not available through Mise with `dnf install -y ...`.
114
+
115
+ `git` and `gh` are not configured inside the container. Do not attempt to use git or gh directly. When the user wants to commit, push, or open a PR, instruct them to `cd` in the worktree directory (matches on host and guest) and run the commands.
@@ -0,0 +1,138 @@
1
+ import contextlib
2
+ import os
3
+ import subprocess
4
+ import sys
5
+ import time
6
+ from collections.abc import AsyncIterator
7
+ from contextlib import asynccontextmanager
8
+ from contextvars import ContextVar
9
+ from io import BytesIO
10
+
11
+ import anyio
12
+ import anyio.abc
13
+ from anyio import create_task_group
14
+ from anyio.abc import ByteReceiveStream, TaskGroup
15
+ from rich.console import Capture
16
+ from rich.text import Text
17
+
18
+ from locki.console import console, err_console
19
+
20
+
21
+ async def _receive_stream(stream: ByteReceiveStream, buffer: BytesIO):
22
+ async for chunk in stream:
23
+ err_console.print(Text.from_ansi(chunk.decode(errors="replace")), style="dim")
24
+ buffer.write(chunk)
25
+
26
+
27
+ @asynccontextmanager
28
+ async def capture_output(
29
+ process: anyio.abc.Process,
30
+ stdout_buf: BytesIO,
31
+ stderr_buf: BytesIO,
32
+ ) -> AsyncIterator[TaskGroup]:
33
+ async with create_task_group() as tg:
34
+ if process.stdout:
35
+ tg.start_soon(_receive_stream, process.stdout, stdout_buf)
36
+ if process.stderr:
37
+ tg.start_soon(_receive_stream, process.stderr, stderr_buf)
38
+ yield tg
39
+
40
+
41
+ async def run_command(
42
+ command: list[str],
43
+ message: str,
44
+ env: dict[str, str] | None = None,
45
+ cwd: str = ".",
46
+ check: bool = True,
47
+ input: bytes | None = None,
48
+ ) -> subprocess.CompletedProcess[bytes]:
49
+ env = env or {}
50
+ try:
51
+ with status(message):
52
+ err_console.print(f"Command: {command}", style="dim")
53
+ start_time = time.time()
54
+ async with await anyio.open_process(
55
+ command,
56
+ stdin=subprocess.PIPE if input else subprocess.DEVNULL,
57
+ env={**os.environ, **env},
58
+ cwd=cwd,
59
+ ) as process:
60
+ stdout_buf, stderr_buf = BytesIO(), BytesIO()
61
+ async with capture_output(process, stdout_buf, stderr_buf):
62
+ if process.stdin and input:
63
+ await process.stdin.send(input)
64
+ await process.stdin.aclose()
65
+ await process.wait()
66
+
67
+ if check and process.returncode != 0:
68
+ raise subprocess.CalledProcessError(
69
+ process.returncode or 0,
70
+ command,
71
+ stdout_buf.getvalue(),
72
+ stderr_buf.getvalue(),
73
+ )
74
+
75
+ elapsed = int(time.time() - start_time)
76
+ duration = (
77
+ "" if elapsed < 5 else f"({elapsed}s)" if elapsed < 60 else f"({elapsed // 60}m{elapsed % 60}s)"
78
+ )
79
+
80
+ if SHOW_SUCCESS_STATUS.get():
81
+ console.print(f"{message} [[green]DONE[/green]] [dim]{duration}[/dim]")
82
+ return subprocess.CompletedProcess(
83
+ command, process.returncode or 0, stdout_buf.getvalue(), stderr_buf.getvalue()
84
+ )
85
+ except FileNotFoundError:
86
+ console.print(f"{message} [[red]ERROR[/red]]")
87
+ console.error(f"{command[0]} is not installed. Please install it first.")
88
+ sys.exit(1)
89
+ except subprocess.CalledProcessError as e:
90
+ console.print(f"{message} [[red]ERROR[/red]]")
91
+ err_console.print(f"[red]Exit code: {e.returncode}[/red]")
92
+ if e.stderr:
93
+ err_console.print(f"[red]Stderr: {e.stderr.decode(errors='replace').strip()}[/red]")
94
+ raise
95
+
96
+
97
+ IN_VERBOSITY_CONTEXT: ContextVar[bool] = ContextVar("in_verbosity_context", default=False)
98
+ VERBOSE: ContextVar[bool] = ContextVar("verbose", default=False)
99
+ SHOW_SUCCESS_STATUS: ContextVar[bool] = ContextVar("show_success_status", default=True)
100
+
101
+
102
+ @contextlib.contextmanager
103
+ def status(message: str):
104
+ if VERBOSE.get():
105
+ console.print(f"{message}...")
106
+ yield
107
+ elif SHOW_SUCCESS_STATUS.get():
108
+ err_console.print(f"\n[bold]{message}[/bold]")
109
+ with console.status(f"{message}...", spinner="dots"):
110
+ yield
111
+ else:
112
+ err_console.print(f"\n[bold]{message}[/bold]")
113
+ yield
114
+
115
+
116
+ @contextlib.contextmanager
117
+ def verbosity(verbose: bool, show_success_status: bool = True):
118
+ if IN_VERBOSITY_CONTEXT.get():
119
+ yield
120
+ return
121
+
122
+ IN_VERBOSITY_CONTEXT.set(True)
123
+ token_verbose = VERBOSE.set(verbose)
124
+ token_status = SHOW_SUCCESS_STATUS.set(show_success_status)
125
+ capture: Capture | None = None
126
+ try:
127
+ with err_console.capture() if not verbose else contextlib.nullcontext() as capture:
128
+ yield
129
+ except Exception:
130
+ if not verbose and capture and (logs := capture.get().strip()):
131
+ err_console.print("\n[yellow]--- Captured logs ---[/yellow]\n")
132
+ err_console.print(Text.from_ansi(logs, style="dim"))
133
+ err_console.print("\n[red]------- Error -------[/red]\n")
134
+ raise
135
+ finally:
136
+ VERBOSE.reset(token_verbose)
137
+ IN_VERBOSITY_CONTEXT.set(False)
138
+ SHOW_SUCCESS_STATUS.reset(token_status)