xas-test-library/console/artifacts.py

"""Bounded log following and result snapshots; never modify original outputs."""
import hashlib
import codecs
import os
from pathlib import Path
import time
from .catalog import safe_path, SECRET


def files_under(paths, limit=256):
    items = []
    deadline = time.monotonic() + 2
    for root in map(Path, paths):
        if root.is_symlink():
            continue
        if root.is_file():
            items.append(root)
        elif root.is_dir():
            for directory, dirs, files in os.walk(root, followlinks=False):
                dirs[:] = sorted(d for d in dirs if not (Path(directory) / d).is_symlink())
                for name in sorted(files):
                    path = Path(directory) / name
                    if not path.is_symlink() and path.is_file():
                        items.append(path)
                    if len(items) >= limit or time.monotonic() > deadline:
                        return items[:limit]
    return items[:limit]


def signature(path):
    try:
        stat = path.stat()
        return (stat.st_ino, stat.st_size, stat.st_mtime_ns)
    except OSError:
        return None


class Collector:
    def __init__(self, settings, directory, output_paths, redactor=None, stream_filter_factory=None, file_limit=256):
        self.settings, self.directory = settings, directory
        self.file_limit = file_limit
        self.output_provider = output_paths if callable(output_paths) else lambda: [Path(p) for p in output_paths]
        self.redactor = redactor
        self.stream_filter_factory = stream_filter_factory
        self.log_filters = {}
        self.before = {str(p): signature(p) for p in files_under(self.outputs, self.file_limit)}
        self.positions = {str(p): p.stat().st_size for p in files_under(self.outputs, self.file_limit) if p.suffix == ".log"}
        self.notices = []
        self.decoders = {}
        self.last_file = None

    @property
    def outputs(self):
        return self.output_provider()

    def follow(self, emit):
        total = 0
        for path in files_under([*self.outputs, self.directory / "work"], self.file_limit):
            if path.suffix != ".log":
                continue
            try:
                start = self.positions.get(str(path), 0)
                if path.stat().st_size < start:
                    start = 0
                    self.decoders.pop(str(path), None)
                with path.open("rb") as handle:
                    handle.seek(start)
                    data = handle.read(32768)
                    self.positions[str(path)] = handle.tell()
                if data:
                    total += len(data)
                    if self.stream_filter_factory:
                        log_filter = self.log_filters.setdefault(str(path), self.stream_filter_factory())
                        data = log_filter.feed(data)
                    decoder = self.decoders.setdefault(str(path), codecs.getincrementaldecoder("utf-8")("replace"))
                    text = decoder.decode(data)
                    prefix = f"\n[文件日志 {path.name}]\n" if self.last_file != str(path) else ""
                    emit((prefix + text).encode("utf-8"))
                    self.last_file = str(path)
            except OSError:
                self.notices.append(f"日志读取失败：{path.name}")
        return total

    def finish_follow(self, emit):
        total = 0
        while total < self.settings.max_log:
            count = self.follow(emit)
            if count == 0:
                break
            total += count
        for name, decoder in self.decoders.items():
            extra = self.log_filters[name].finish() if name in self.log_filters else b""
            tail = decoder.decode(extra, final=True)
            if tail:
                emit(tail.encode("utf-8"))
        self.decoders.clear()
        self.log_filters.clear()

    def external_growth(self):
        growth = 0
        for path in files_under(self.outputs, limit=1025):
            stat = signature(path)
            before = self.before.get(str(path))
            if stat:
                growth += max(0, stat[1] - (before[1] if before and before[0] == stat[0] else 0))
        return growth

    def collect(self):
        destination = self.directory / "artifacts"
        destination.mkdir(exist_ok=True)
        local = files_under([self.directory / "work"])
        changed = [p for p in files_under(self.outputs, self.file_limit) if signature(p) != self.before.get(str(p))]
        sources = [self.directory / "source" / "script.txt", *local, *changed]
        items, total = [], 0
        for path in sources[:self.settings.max_run_files - 1]:
            try:
                if not path.is_file() or path.is_symlink():
                    continue
                if path.stat().st_size > self.settings.max_artifact:
                    retention = "临时文件将在收尾时释放" if path in local else "原位置文件保留"
                    self.notices.append(f"未归档 {path.name}：超过单文件 32 MiB 上限，{retention}")
                    continue
                data = path.read_bytes()
                if len(data) > self.settings.max_artifact or total + len(data) > 64 * 1024 * 1024:
                    self.notices.append("结果超过归档容量限制：原目录文件保留，未归档的临时结果将在收尾时释放")
                    continue
                suffix = path.suffix.lower()
                kind = "json" if suffix == ".json" else ("text" if suffix in {".log", ".txt", ".py", ".sh", ".csv", ".md", ".jsonl"} else "file")
                if kind != "file":
                    text = data.decode("utf-8", errors="replace")
                    data = SECRET.sub("[已隐藏疑似凭据]", text).encode("utf-8")
                    if self.redactor:
                        data = self.redactor(data)
                identity = hashlib.sha256(str(path).encode()).hexdigest()[:24]
                target = destination / identity
                with target.open("xb") as handle:
                    handle.write(data)
                items.append({"id": identity, "name": path.name, "size": len(data), "kind": kind})
                total += len(data)
            except OSError:
                self.notices.append(f"未能归档：{path.name}")
        if len(sources) > self.settings.max_run_files - 1:
            self.notices.append("结果超过文件数上限：保留 63 个文件及运行摘要，未归档临时结果将释放")
        return items