from __future__ import annotations import ast import hashlib import os import stat import tempfile from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path from uuid import uuid4 from packaging.requirements import InvalidRequirement, Requirement from packaging.utils import canonicalize_name from .validation import validate_distribution_name, validate_import_name class EditorError(RuntimeError): pass @dataclass(frozen=True, slots=True) class MutationReceipt: path: Path changed: bool backup_path: Path | None existed_before: bool previous_text: str before_sha256: str after_sha256: str class AtomicTextEditor: def __init__(self, path: Path, backup_dir: Path, backups_to_keep: int): raw_path = Path(path).expanduser() if raw_path.exists(): if not raw_path.is_file(): raise EditorError(f"Managed path is not a regular file: {raw_path}") self.path = raw_path.resolve(strict=True) else: self.path = raw_path.parent.resolve(strict=True) / raw_path.name self.backup_dir = Path(backup_dir) self.backups_to_keep = backups_to_keep def read(self, *, required: bool = False) -> str: if not self.path.exists(): if required: raise EditorError(f"Managed file does not exist: {self.path}") return "" try: return self.path.read_text(encoding="utf-8") except (OSError, UnicodeDecodeError) as exc: raise EditorError(f"Unable to read managed file: {self.path.name}") from exc def apply(self, new_text: str, *, required: bool = False) -> MutationReceipt: old_text = self.read(required=required) existed = self.path.exists() before = hashlib.sha256(old_text.encode()).hexdigest() after = hashlib.sha256(new_text.encode()).hexdigest() if old_text == new_text: return MutationReceipt(self.path, False, None, existed, old_text, before, after) backup = self._backup(old_text) if existed else None self._atomic_write(new_text) self._prune_backups() return MutationReceipt(self.path, True, backup, existed, old_text, before, after) def rollback(self, receipt: MutationReceipt) -> None: if not receipt.changed: return if receipt.existed_before: self._atomic_write(receipt.previous_text) else: self.path.unlink(missing_ok=True) def _backup(self, text: str) -> Path: self.backup_dir.mkdir(mode=0o700, parents=True, exist_ok=True) stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ") backup = self.backup_dir / f"{self.path.name}.{stamp}.{uuid4().hex}.bak" try: with backup.open("x", encoding="utf-8", newline="") as handle: handle.write(text) handle.flush() os.fsync(handle.fileno()) os.chmod(backup, 0o600) except OSError as exc: backup.unlink(missing_ok=True) raise EditorError(f"Unable to create backup for {self.path.name}") from exc return backup def _atomic_write(self, text: str) -> None: existing_stat = self.path.stat() if self.path.exists() else None mode = stat.S_IMODE(existing_stat.st_mode) if existing_stat else 0o640 fd, temp_name = tempfile.mkstemp(prefix=f".{self.path.name}.", dir=self.path.parent) temp_path = Path(temp_name) try: with os.fdopen(fd, "w", encoding="utf-8", newline="") as handle: handle.write(text) handle.flush() os.fsync(handle.fileno()) os.chmod(temp_path, mode) if existing_stat is not None and hasattr(os, "chown"): try: os.chown(temp_path, existing_stat.st_uid, existing_stat.st_gid) except PermissionError: pass os.replace(temp_path, self.path) if os.name != "nt": directory_fd = os.open(self.path.parent, os.O_RDONLY) try: os.fsync(directory_fd) finally: os.close(directory_fd) except OSError as exc: temp_path.unlink(missing_ok=True) raise EditorError(f"Unable to atomically update {self.path.name}") from exc def _prune_backups(self) -> None: pattern = f"{self.path.name}.*.bak" backups = sorted(self.backup_dir.glob(pattern), key=lambda path: path.stat().st_mtime, reverse=True) for old_backup in backups[self.backups_to_keep :]: old_backup.unlink(missing_ok=True) class PluginConfigurationEditor(AtomicTextEditor): def enabled_plugins(self) -> list[str]: text = self.read(required=True) _, plugins = self._find_plugins_assignment(text) return plugins def preview(self, import_name: str, enabled: bool) -> tuple[str, list[str], bool]: import_name = validate_import_name(import_name) text = self.read(required=True) node, plugins = self._find_plugins_assignment(text) changed = False if enabled and import_name not in plugins: plugins.append(import_name) changed = True elif not enabled and import_name in plugins: plugins = [plugin for plugin in plugins if plugin != import_name] changed = True if not changed: return text, plugins, False newline = "\r\n" if "\r\n" in text else "\n" lines = text.splitlines(keepends=True) if node.lineno == node.end_lineno and ";" in lines[node.lineno - 1]: raise EditorError("PLUGINS assignment sharing a line cannot be safely edited.") replacement = [f"PLUGINS = [{newline}"] replacement.extend(f" {plugin!r},{newline}" for plugin in plugins) replacement.append(f"]{newline}") new_text = "".join(lines[: node.lineno - 1] + replacement + lines[node.end_lineno :]) return new_text, plugins, True def set_enabled(self, import_name: str, enabled: bool) -> MutationReceipt: new_text, _, _ = self.preview(import_name, enabled) return self.apply(new_text, required=True) @staticmethod def _find_plugins_assignment(text: str) -> tuple[ast.Assign | ast.AnnAssign, list[str]]: try: tree = ast.parse(text) except SyntaxError as exc: raise EditorError("configuration.py is not valid Python; refusing to edit it.") from exc matches: list[ast.Assign | ast.AnnAssign] = [] for node in tree.body: if isinstance(node, ast.Assign) and any( isinstance(target, ast.Name) and target.id == "PLUGINS" for target in node.targets ): matches.append(node) elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name) and node.target.id == "PLUGINS": matches.append(node) if len(matches) != 1: raise EditorError("configuration.py must contain exactly one static PLUGINS assignment.") node = matches[0] try: value = ast.literal_eval(node.value) except (ValueError, TypeError, SyntaxError) as exc: raise EditorError("PLUGINS must be a literal list or tuple of import names.") from exc if not isinstance(value, (list, tuple)) or any(not isinstance(item, str) for item in value): raise EditorError("PLUGINS must be a literal list or tuple of import names.") plugins = [validate_import_name(item) for item in value] if len(set(plugins)) != len(plugins): raise EditorError("PLUGINS contains duplicate entries; refusing to edit it.") return node, plugins class RequirementsEditor(AtomicTextEditor): def preview(self, package_name: str, requirement_line: str | None) -> tuple[str, bool]: package_name = validate_distribution_name(package_name) wanted = canonicalize_name(package_name) text = self.read(required=False) lines = text.splitlines(keepends=True) indexes: list[int] = [] for index, line in enumerate(lines): stripped = line.strip() if not stripped or stripped.startswith("#") or stripped.startswith(("-r", "--", "-e")): continue candidate = stripped.split(" #", 1)[0].rstrip() try: requirement = Requirement(candidate) except InvalidRequirement: continue if canonicalize_name(requirement.name) == wanted: indexes.append(index) if len(indexes) > 1: raise EditorError("local_requirements.txt contains duplicate entries for this package.") newline = "\r\n" if "\r\n" in text else "\n" replacement = f"{requirement_line}{newline}" if requirement_line else None if requirement_line: try: parsed = Requirement(requirement_line) except InvalidRequirement as exc: raise EditorError("Generated requirement is invalid.") from exc if canonicalize_name(parsed.name) != wanted: raise EditorError("Generated requirement targets the wrong distribution.") if indexes: index = indexes[0] if replacement is None: del lines[index] elif lines[index].rstrip("\r\n") == requirement_line: return text, False else: lines[index] = replacement elif replacement is not None: if text and not text.endswith(("\n", "\r")): lines.append(newline) lines.append(replacement) else: return text, False return "".join(lines), True def set_requirement(self, package_name: str, requirement_line: str | None) -> MutationReceipt: new_text, _ = self.preview(package_name, requirement_line) return self.apply(new_text, required=False)