93 lines
3.1 KiB
Python
93 lines
3.1 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import subprocess
|
||
|
|
import os
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from tools.capabilities import Capability
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class ToolResult:
|
||
|
|
returncode: int
|
||
|
|
stdout: str
|
||
|
|
stderr: str
|
||
|
|
|
||
|
|
|
||
|
|
class CapabilityError(PermissionError):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class WorktreeTools:
|
||
|
|
def __init__(self, worktree_path: Path, capabilities: set[Capability]) -> None:
|
||
|
|
self.worktree_path = worktree_path.resolve()
|
||
|
|
self.capabilities = capabilities
|
||
|
|
|
||
|
|
def _require(self, capability: Capability) -> None:
|
||
|
|
if capability not in self.capabilities:
|
||
|
|
raise CapabilityError(f"Missing capability: {capability.value}")
|
||
|
|
|
||
|
|
def _safe_path(self, relative_path: str) -> Path:
|
||
|
|
path = (self.worktree_path / relative_path).resolve()
|
||
|
|
if self.worktree_path not in path.parents and path != self.worktree_path:
|
||
|
|
raise CapabilityError("Path escapes worktree")
|
||
|
|
return path
|
||
|
|
|
||
|
|
def read_text(self, relative_path: str) -> str:
|
||
|
|
self._require(Capability.READ_REPOSITORY)
|
||
|
|
return self._safe_path(relative_path).read_text(encoding="utf-8")
|
||
|
|
|
||
|
|
def write_text(self, relative_path: str, content: str) -> None:
|
||
|
|
self._require(Capability.WRITE_WORKTREE)
|
||
|
|
path = self._safe_path(relative_path)
|
||
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
path.write_text(content, encoding="utf-8")
|
||
|
|
|
||
|
|
def run(self, command: list[str], timeout: int = 60) -> ToolResult:
|
||
|
|
self._require(Capability.RUN_TESTS)
|
||
|
|
env = os.environ.copy()
|
||
|
|
for key in ["DJANGO_SETTINGS_MODULE", "DATABASE_URL"]:
|
||
|
|
env.pop(key, None)
|
||
|
|
completed = subprocess.run(
|
||
|
|
command,
|
||
|
|
cwd=self.worktree_path,
|
||
|
|
check=False,
|
||
|
|
capture_output=True,
|
||
|
|
text=True,
|
||
|
|
timeout=timeout,
|
||
|
|
env=env,
|
||
|
|
)
|
||
|
|
return ToolResult(completed.returncode, completed.stdout, completed.stderr)
|
||
|
|
|
||
|
|
def git(self, args: list[str]) -> ToolResult:
|
||
|
|
required = Capability.COMMIT_CHANGES if args and args[0] == "commit" else Capability.READ_REPOSITORY
|
||
|
|
self._require(required)
|
||
|
|
completed = subprocess.run(
|
||
|
|
["git", *args],
|
||
|
|
cwd=self.worktree_path,
|
||
|
|
check=False,
|
||
|
|
capture_output=True,
|
||
|
|
text=True,
|
||
|
|
)
|
||
|
|
return ToolResult(completed.returncode, completed.stdout, completed.stderr)
|
||
|
|
|
||
|
|
def status(self) -> str:
|
||
|
|
return self.git(["status", "--short"]).stdout
|
||
|
|
|
||
|
|
def diff(self) -> str:
|
||
|
|
return self.git(["diff", "--", "."]).stdout
|
||
|
|
|
||
|
|
def commit_all(self, message: str) -> str:
|
||
|
|
self._require(Capability.COMMIT_CHANGES)
|
||
|
|
add = self.git(["add", "."])
|
||
|
|
if add.returncode != 0:
|
||
|
|
raise RuntimeError(add.stderr)
|
||
|
|
commit = self.git(["-c", "user.name=Artifex", "-c", "user.email=artifex@example.invalid", "commit", "-m", message])
|
||
|
|
if commit.returncode != 0:
|
||
|
|
raise RuntimeError(commit.stderr or commit.stdout)
|
||
|
|
rev = self.git(["rev-parse", "HEAD"])
|
||
|
|
if rev.returncode != 0:
|
||
|
|
raise RuntimeError(rev.stderr)
|
||
|
|
return rev.stdout.strip()
|