238 lines
9.9 KiB
Python
238 lines
9.9 KiB
Python
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)
|