Something went wrong. Try again.
Something went wrong. Try again.
22 kB · 491 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492"""Bounded reads, exact edits, and ripgrep-backed search over workspace files.
Every external command runs through the supervised `bash` job plugin, so thismodule never spawns a process of its own and every child stays in a job's group."""from __future__ import annotations
import difflibimport fnmatchimport jsonimport osimport reimport shleximport statimport tempfilefrom dataclasses import dataclassfrom pathlib import Pathfrom collections.abc import Awaitable, Callable, Generatorfrom typing import Generic, Iterable, Sequence, TypeVar
from albedo_api import PythonApi, ReadyList, Textfrom albedo_plugins import bash as jobs
READ_LIMIT = 16_000SEARCH_TIMEOUT = 30LINE_WIDTH = 160CANDIDATES = 4LISTED = 10CONTEXT = 2PREVIEW_LINES = 2
@dataclass(frozen=True)class Match: """One matching line, or with `context=` a surrounding line that prints grep-style as `path-line- text`.""" path: str line: int text: str context: bool = False
def __str__(self) -> str: mark = "-" if self.context else ":" return f"{self.path}{mark}{self.line}{mark} {self.text}"
def to_dict(self) -> dict[str, str | int | bool]: fields: dict[str, str | int | bool] = {"path": self.path, "line": self.line, "text": self.text} if self.context: fields["context"] = True return fields
class Rows(ReadyList): """A list that prints one row per line, so a REPL result reads like output. Like Text, it may be awaited or used directly."""
def __init__(self, items: Iterable[object] = (), truncated: bool = False) -> None: super().__init__(items) self.truncated = truncated
def __repr__(self) -> str: if not self: return "[]" body = "\n".join(str(item) for item in self) return body + ("\n[truncated]" if self.truncated else "")
__str__ = __repr__
Result = TypeVar("Result")
class Search(Generic[Result]): """Work that runs a supervised job, so its result exists only once awaited: `Search[Rows]` awaits to Rows. Using it without `await` explains that instead of printing a coroutine."""
def __init__(self, call: str, run: Callable[[], Awaitable[Result]], result: str = "rows") -> None: self._call = call self._run = run self._result = result
def __await__(self) -> Generator[object, None, Result]: return self._run().__await__()
def __getitem__(self, index: object) -> "Search[Result]": """`await files.find(...)[:10]` slices before it awaits; apply the slice to the rows instead of failing on precedence.""" async def run() -> object: rows = await self._run() picked = rows[index] # type: ignore[index] if isinstance(rows, Rows) and isinstance(index, slice): return Rows(picked, truncated=rows.truncated) return picked return Search(self._call, run, self._result) # type: ignore[arg-type]
def _unawaited(self) -> TypeError: return TypeError(f"{self._call}(...) runs in the background; " f"use `await {self._call}(...)` to get its {self._result}")
def __repr__(self) -> str: return f"<{self._call}(...) has not run: use `await {self._call}(...)` for its {self._result}>"
__str__ = __repr__
def __iter__(self): raise self._unawaited()
def __len__(self) -> int: raise self._unawaited()
async def _run(command: str) -> tuple[int | None, str]: """One supervised shell job, awaited to completion, with its bounded output.""" job = jobs.bash(command, timeout=SEARCH_TIMEOUT) try: await job return job.exit_code, job.tail(jobs.preview_limit) finally: # The search is the plugin's own work, not a job the model started. if job.exit_code is not None or job.timed_out: jobs.forget(job)
def _numbered(number: int, lines: Sequence[str]) -> str: text = lines[number - 1] if 0 < number <= len(lines) else "" return f"{number:>6} | {text[:LINE_WIDTH]}"
class Files: """Workspace file access with the diagnostics an exact edit needs."""
def read(self, path: str, start_line: int = 1, end_line: int | None = None, *, limit: int | None = None, max_chars: int = READ_LIMIT) -> Text: """Numbered lines from start_line through end_line, at most `limit` lines.
`limit` counts lines. `max_chars` is a separate character budget (at most 200000) that keeps a huge window from flooding the context; when it stops a read early, the result says so and names the line to resume from. Lines are never shortened. The numbers are what `edit(line_hint=)` takes. """ if start_line < 1 or (end_line is not None and end_line < start_line) \ or (limit is not None and limit < 1) or not 0 < max_chars <= 200_000: raise ValueError("start_line >= 1, end_line >= start_line, limit >= 1 line, " "0 < max_chars <= 200000") target = Path(path).expanduser() if not target.is_file(): if target.is_dir(): raise IsADirectoryError(f"{path} is a directory; files.ls({path!r}) lists it") raise FileNotFoundError(_missing(target, path)) lines = target.read_bytes().decode("utf-8", errors="replace").splitlines() if start_line > len(lines): return Text(f"[{path} has {len(lines)} lines; nothing at line {start_line}]") requested = len(lines) if end_line is None else min(end_line, len(lines)) last = requested if limit is None else min(requested, start_line + limit - 1) body, used = [], 0 for number in range(start_line, last + 1): row = f"{number:>6} | {lines[number - 1]}" if used + len(row) + 1 > max_chars: if not body: retry = ( f"read it alone with start_line={number}, end_line={number}, max_chars={len(row) + 1}" if len(row) + 1 <= 200_000 else "it exceeds the 200000-character maximum" ) body.append(f"[line {number} is {len(row)} characters with its number, over " f"max_chars={max_chars}; {retry}]") else: body.append(f"[stopped at max_chars={max_chars} characters; lines {number}-{last} " f"not shown; read again with start_line={number}, or raise max_chars]") break used += len(row) + 1 body.append(row) else: if last < requested: body.append(f"[limit={limit} lines reached; {requested - last} more through line " f"{requested}; read again with start_line={last + 1}]") return Text("\n".join(body))
def ls(self, path: str = ".", pattern: str | None = None, *, hidden: bool = False) -> Rows: """One directory, directories suffixed with `/`.""" directory = Path(path).expanduser() entries = [] for entry in sorted(directory.iterdir(), key=lambda item: item.name): if not hidden and entry.name.startswith("."): continue if pattern and not fnmatch.fnmatch(entry.name, pattern): continue entries.append(entry.name + ("/" if entry.is_dir() else "")) return Rows(entries)
def find(self, pattern: str, path: str | Sequence[str] = ".", *, glob: str | Sequence[str] | None = None, context: int = 0, max_results: int = 50, literal: bool = False, case_sensitive: bool | None = None, hidden: bool = False) -> Search[Rows]: """Content search through ripgrep when it is installed, else pure Python. `path` may be one path or a list; `context=N` adds N lines around each match. Await it: `await files.find(pattern)`.""" if not 0 <= context <= 50: raise ValueError("0 <= context <= 50") return Search("files.find", lambda: self._find(pattern, path, glob, context, max_results, literal, case_sensitive, hidden))
async def _find(self, pattern: str, path: str | Sequence[str], glob: str | Sequence[str] | None, context: int, max_results: int, literal: bool, case_sensitive: bool | None, hidden: bool) -> Rows: targets = [Path(item).expanduser() for item in ([path] if isinstance(path, str) else path)] if not _which("rg"): return _fallback_find(pattern, targets, glob, context, max_results, literal, case_sensitive, hidden) flags = ["--json"] if context: flags += ["-C", str(context)] if literal: flags.append("-F") if case_sensitive is True: flags.append("-s") elif case_sensitive is False: flags.append("-i") if hidden: flags.append("--hidden") for value in ([glob] if isinstance(glob, str) else list(glob or [])): flags += ["-g", value] command = " ".join(shlex.quote(part) for part in ["rg", *flags, "-e", pattern, *map(str, targets)]) per_match = 4 + 2 * context status, output = await _run(f"{command} | head -n {max(1, max_results) * per_match}") if status not in (0, 1, None) and not output.strip(): raise RuntimeError(f"search failed: {output.strip()[:400]}") results, matched = [], 0 for line in output.splitlines(): try: event = json.loads(line) except ValueError: continue kind = event.get("type") if kind not in ("match", "context"): continue if kind == "match" and matched >= max_results: return Rows(results, truncated=True) data = event["data"] results.append(Match(data["path"]["text"], data["line_number"], data["lines"]["text"].rstrip("\r\n")[:LINE_WIDTH * 4], context=kind == "context")) matched += kind == "match" return Rows(results)
def paths(self, pattern: str | None = None, path: str = ".", *, glob: str | None = None, max_results: int = 100, hidden: bool = False) -> Search[Rows]: """File names, not contents; the same ripgrep-or-Python split. A pattern with *, ? or [ is a glob over names; other text matches anywhere in the path. Await it: `await files.paths(pattern)`.""" return Search("files.paths", lambda: self._paths(pattern, path, glob, max_results, hidden))
async def _paths(self, pattern: str | None, path: str, glob: str | None, max_results: int, hidden: bool) -> Rows: target = Path(path).expanduser() if not _which("rg"): return _fallback_paths(pattern, target, glob, max_results, hidden) flags = ["--files"] if hidden: flags.append("--hidden") if glob: flags += ["-g", glob] command = " ".join(shlex.quote(part) for part in ["rg", *flags, str(target)]) _, output = await _run(command) matches = _name_matcher(pattern) found = [line for line in output.splitlines() if line and matches(line)] return Rows(found[:max_results], truncated=len(found) > max_results)
def edit(self, path: str, old_str: str, new_str: str, line_hint: int | None = None) -> Text: """Replace one exact, unique string. A miss reports what is actually there.""" if not old_str: raise ValueError("old_str must be non-empty") if line_hint is not None and (isinstance(line_hint, bool) or not isinstance(line_hint, int)): raise ValueError(f"line_hint must be a line number, not {type(line_hint).__name__}") target = Path(path).expanduser() if not target.exists(): raise FileNotFoundError(_missing(target, path)) target = target.resolve() snapshot = _snapshot(target) content = snapshot[0].decode("utf-8") found = _occurrences(content, old_str) if not found: closest = _closest(content, old_str) raise ValueError(f"string not found in {path}" + (f"\n{closest}" if closest else "")) chosen = _choose(content, old_str, found, line_hint, path) replacement = (content[:chosen.index] + new_str + content[chosen.index + len(old_str):]).encode("utf-8") _replace(target, snapshot, replacement) return Text(f"Edited {target}")
def write(self, path: str, content: str) -> Text: """Create or replace a whole file. Use `edit` to change part of one.""" target = Path(path).expanduser() target.parent.mkdir(parents=True, exist_ok=True) target.write_text(content, encoding="utf-8") return Text(f"Wrote {target.resolve()} ({len(content.encode())} bytes)")
def _which(name: str) -> str | None: for directory in os.environ.get("PATH", "").split(os.pathsep): candidate = Path(directory)/name if candidate.is_file() and os.access(candidate, os.X_OK): return str(candidate) return None
def _fallback_find(pattern, targets: list[Path], glob, context, max_results, literal, case_sensitive, hidden) -> Rows: flags = 0 if case_sensitive else re.IGNORECASE matcher = re.compile(re.escape(pattern) if literal else pattern, flags) globs = [glob] if isinstance(glob, str) else list(glob or []) results, matched = [], 0 for target in targets: for file in _walk(target, hidden): if globs and not any(fnmatch.fnmatch(str(file), value) or fnmatch.fnmatch(file.name, value) for value in globs): continue try: lines = file.read_text(encoding="utf-8", errors="replace").splitlines() except OSError: continue hits = [number for number, line in enumerate(lines, 1) if matcher.search(line)] hit = set(hits) shown = 0 for number in hits: if matched >= max_results: return Rows(results, truncated=True) for around in range(max(number - context, shown + 1), min(number + context, len(lines)) + 1): results.append(Match(str(file), around, lines[around - 1][:LINE_WIDTH * 4], context=around not in hit)) shown = around matched += 1 return Rows(results)
def _name_matcher(pattern: str | None) -> Callable[[str], bool]: """A pattern with *, ? or [ is a glob over the file name (or the whole path when it has a /); anything else matches as case-insensitive text anywhere in the path, so "*.md" and "readme" both mean what they look like.""" if not pattern: return lambda _path: True folded = pattern.lower() if any(char in pattern for char in "*?["): return lambda path: fnmatch.fnmatch((path if "/" in pattern else Path(path).name).lower(), folded) return lambda path: folded in path.lower()
def _fallback_paths(pattern, path: Path, glob, max_results, hidden) -> Rows: matches = _name_matcher(pattern) found = [str(file) for file in _walk(path, hidden) if (glob is None or fnmatch.fnmatch(file.name, glob)) and matches(str(file))] return Rows(found[:max_results], truncated=len(found) > max_results)
def _walk(path: Path, hidden: bool): if path.is_file(): yield path return for root, directories, names in os.walk(path): if not hidden: directories[:] = [name for name in directories if not name.startswith(".")] names = [name for name in names if not name.startswith(".")] for name in sorted(names): yield Path(root)/name
def _snapshot(target: Path) -> tuple[bytes, tuple[int, int, int, int, int]]: """Bytes plus identity: an edit refuses to publish over a concurrent change.""" for _ in range(3): before = target.stat() data = target.read_bytes() after = target.stat() identity = (after.st_dev, after.st_ino, stat.S_IMODE(after.st_mode), after.st_size, after.st_mtime_ns) if identity == (before.st_dev, before.st_ino, stat.S_IMODE(before.st_mode), before.st_size, before.st_mtime_ns): return data, identity raise ValueError(f"file changed while reading {target}; refusing to edit it")
def _replace(target: Path, snapshot: tuple[bytes, tuple[int, int, int, int, int]], replacement: bytes) -> None: descriptor, name = tempfile.mkstemp(dir=target.parent, prefix=f".{target.name}.", suffix=".tmp") temporary = Path(name) try: with os.fdopen(descriptor, "wb") as stream: os.fchmod(stream.fileno(), snapshot[1][2]) stream.write(replacement) stream.flush() os.fsync(stream.fileno()) if _snapshot(target) != snapshot: raise ValueError(f"file changed while editing {target}; refusing to overwrite concurrent changes") os.replace(temporary, target) finally: try: temporary.unlink() except FileNotFoundError: pass
@dataclass(frozen=True)class _Occurrence: index: int start_line: int end_line: int
@property def label(self) -> str: return str(self.start_line) if self.start_line == self.end_line else f"{self.start_line}-{self.end_line}"
def _occurrences(content: str, needle: str) -> list[_Occurrence]: found, span, start = [], needle.count("\n") + 1, 0 while True: index = content.find(needle, start) if index < 0: return found first = content.count("\n", 0, index) + 1 found.append(_Occurrence(index, first, first + span - 1)) start = index + len(needle)
def _choose(content: str, old_str: str, found: list[_Occurrence], line_hint: int | None, path: str) -> _Occurrence: if len(found) == 1: return found[0] listed = ", ".join(one.label for one in found[:LISTED]) + (f", ... ({len(found)} total)" if len(found) > LISTED else "") if line_hint is None: raise ValueError( f"found {len(found)} occurrences in {path}: lines {listed}. Nothing was changed. " "Retry with line_hint=<a line inside the range you want>, or widen old_str until it is unique.\n" + _candidates(content, found)) inside = next((one for one in found if one.start_line <= line_hint <= one.end_line), None) if inside is None: raise ValueError( f"line_hint={line_hint} is inside none of the {len(found)} occurrences in {path}: lines {listed}. " f"Nothing was changed.\n{_candidates(content, found)}") return inside
def _candidates(content: str, found: list[_Occurrence]) -> str: lines = content.splitlines() blocks = [] for position, occurrence in enumerate(found[:CANDIDATES], 1): matched = occurrence.end_line - occurrence.start_line + 1 preview = min(matched, PREVIEW_LINES) first = max(1, occurrence.start_line - CONTEXT) last = min(len(lines), occurrence.end_line + CONTEXT) body = [_numbered(number, lines) for number in range(first, occurrence.start_line)] body += [_numbered(number, lines) for number in range(occurrence.start_line, occurrence.start_line + preview)] if matched > preview: body.append(f"{'...':>6} | ({matched - preview} further matching lines)") body += [_numbered(number, lines) for number in range(occurrence.end_line + 1, last + 1)] blocks.append(f"candidate {position} of {len(found)}: lines {occurrence.label}\n" + "\n".join(body)) if len(found) > CANDIDATES: blocks.append(f"... {len(found) - CANDIDATES} further occurrences not shown") return "\n\n".join(blocks)
def _closest(content: str, needle: str) -> str: lines, needle_lines = content.splitlines(), needle.splitlines() head = next((line.strip() for line in needle_lines if line.strip()), "") if not lines or not head: return "" ranked = sorted(((difflib.SequenceMatcher(None, head, line.strip(), autojunk=False).ratio(), index) for index, line in enumerate(lines) if line.strip()), reverse=True)[:12] width = max(1, len(needle_lines)) scored = [] for _, start in ranked: end = min(len(lines), start + width) current = "\n".join(lines[start:end]).strip() scored.append((difflib.SequenceMatcher(None, needle.strip(), current, autojunk=False).ratio(), start, end)) blocks = [] for score, start, end in sorted(scored, reverse=True): if score < 0.25: continue body = "\n".join(_numbered(number + 1, lines) for number in range(start, end)) blocks.append(f"closest candidate lines {start + 1}-{end} (similarity {score:.0%}):\n{body}") if len(blocks) == 3: break return "\n\n".join(blocks)
def _missing(target: Path, original: str) -> str: message = f"{original} not found (cwd: {Path.cwd()})" parent = target.parent if not parent.is_dir(): return message close = difflib.get_close_matches(target.name, [entry.name for entry in parent.iterdir()], n=3, cutoff=0.35) return message + (f"; nearby paths: {', '.join(str(parent/name) for name in close)}" if close else "")
def setup(api: PythonApi) -> dict[str, object]: return {"files": Files()}