From 04498c7d4239cccf0e0ad2eeb9bd248456fa3385 Mon Sep 17 00:00:00 2001 From: stumpylog <797416+stumpylog@users.noreply.github.com> Date: Wed, 22 Jul 2026 14:57:48 -0700 Subject: [PATCH] Feature: add ExportSink ABC and DirectoryExportSink --- src/documents/export/sinks.py | 172 +++++++++++++++++++++++ src/documents/tests/export/test_sinks.py | 141 +++++++++++++++++++ 2 files changed, 313 insertions(+) diff --git a/src/documents/export/sinks.py b/src/documents/export/sinks.py index 43c7f073e..53eb6f414 100644 --- a/src/documents/export/sinks.py +++ b/src/documents/export/sinks.py @@ -1,11 +1,19 @@ from __future__ import annotations +import hashlib import json +from contextlib import contextmanager from typing import TYPE_CHECKING from django.core.serializers.json import DjangoJSONEncoder +from documents.file_handling import delete_empty_directories +from documents.utils import compute_checksum +from documents.utils import copy_file_with_basic_stats + if TYPE_CHECKING: + from collections.abc import Iterator + from pathlib import Path from typing import TextIO @@ -41,3 +49,167 @@ class StreamingManifestWriter: def close(self) -> None: """Write the closing bracket. Does NOT close the handle (the sink owns it).""" self._file.write("\n]") + + +class ExportSink: + """Destination for a document export. + + The command declares export contents via three verbs; the sink decides how to + persist each. ``arcname`` is always a relative POSIX path + (e.g. ``"manifest.json"``, ``"originals/foo.pdf"``). + + Contract: + * At most one ``stream()`` open at a time (it is the manifest); + ``add_file``/``add_json`` may be called while it is open. + * Context-manager: normal exit finalizes, an exception aborts. No partial or + failed run leaves a complete-looking artifact. + """ + + def add_file( + self, + source: Path, + arcname: str, + *, + checksum: str | None = None, + ) -> None: + raise NotImplementedError # pragma: no cover + + def add_json(self, content: list | dict, arcname: str) -> None: + raise NotImplementedError # pragma: no cover + + def stream(self, arcname: str): # -> contextmanager yielding TextIO + raise NotImplementedError # pragma: no cover + + def _open(self) -> None: + """Hook called on context entry. Override as needed.""" + + def _finalize(self) -> None: + """Commit on clean exit.""" + raise NotImplementedError # pragma: no cover + + def _abort(self) -> None: + """Roll back on exception.""" + raise NotImplementedError # pragma: no cover + + def __enter__(self) -> ExportSink: + self._open() + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + if exc_type is not None: + self._abort() + else: + self._finalize() + + +class DirectoryExportSink(ExportSink): + """Writes loose files into a target directory, with incremental sync. + + Owns the snapshot/skip/compare/prune machinery that used to live in the + command (``files_in_export_dir``, ``check_and_copy``, ``check_and_write_json``, + and the ``--delete`` pass). + """ + + def __init__( + self, + target: Path, + *, + compare_checksums: bool, + compare_json: bool, + delete: bool, + ) -> None: + self._target = target.resolve() + self._compare_checksums = compare_checksums + self._compare_json = compare_json + self._delete = delete + self._snapshot: set[Path] = set() + self._stream_open = False + + def _open(self) -> None: + for x in self._target.glob("**/*"): + if x.is_file(): + self._snapshot.add(x.resolve()) + + def add_file( + self, + source: Path, + arcname: str, + *, + checksum: str | None = None, + ) -> None: + target = (self._target / arcname).resolve() + self._snapshot.discard(target) + perform_copy = False + if target.exists(): + source_stat = source.stat() + target_stat = target.stat() + if self._compare_checksums and checksum: + perform_copy = compute_checksum(target) != checksum + elif ( + source_stat.st_mtime != target_stat.st_mtime + or source_stat.st_size != target_stat.st_size + ): + perform_copy = True + else: + perform_copy = True + if perform_copy: + target.parent.mkdir(parents=True, exist_ok=True) + copy_file_with_basic_stats(source, target) + + def add_json(self, content: list | dict, arcname: str) -> None: + target = (self._target / arcname).resolve() + json_str = _dumps(content) + perform_write = True + if target in self._snapshot: + self._snapshot.discard(target) + if self._compare_json: + target_checksum = hashlib.blake2b(target.read_bytes()).hexdigest() + src_checksum = hashlib.blake2b(json_str.encode("utf-8")).hexdigest() + if src_checksum == target_checksum: + perform_write = False + if perform_write: + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(json_str, encoding="utf-8") + + @contextmanager + def stream(self, arcname: str) -> Iterator[TextIO]: + if self._stream_open: + raise RuntimeError("A stream is already open on this sink") + target = (self._target / arcname).resolve() + tmp = target.with_suffix(target.suffix + ".tmp") + target.parent.mkdir(parents=True, exist_ok=True) + handle = tmp.open("w", encoding="utf-8") + self._stream_open = True + try: + yield handle + except BaseException: + handle.close() + tmp.unlink(missing_ok=True) + raise + else: + handle.close() + self._commit_streamed_file(target, tmp) + finally: + self._stream_open = False + + def _commit_streamed_file(self, target: Path, tmp: Path) -> None: + if target in self._snapshot: + self._snapshot.discard(target) + if self._compare_json: + existing = hashlib.blake2b(target.read_bytes()).hexdigest() + new = hashlib.blake2b(tmp.read_bytes()).hexdigest() + if existing == new: + tmp.unlink() + return + tmp.rename(target) + + def _finalize(self) -> None: + if self._delete: + for f in self._snapshot: + f.unlink() + delete_empty_directories(f.parent, self._target) + + def _abort(self) -> None: + # Folder mode is in-place/incremental: streamed .tmp files are already + # cleaned in stream(); leave everything else intact and skip the prune. + return None diff --git a/src/documents/tests/export/test_sinks.py b/src/documents/tests/export/test_sinks.py index 296622f02..0a6577a57 100644 --- a/src/documents/tests/export/test_sinks.py +++ b/src/documents/tests/export/test_sinks.py @@ -1,6 +1,11 @@ import io import json +import os +from pathlib import Path +import pytest + +from documents.export.sinks import DirectoryExportSink from documents.export.sinks import StreamingManifestWriter from documents.export.sinks import _dumps @@ -27,3 +32,139 @@ class TestStreamingManifestWriter: writer: StreamingManifestWriter = StreamingManifestWriter(handle) writer.close() assert json.loads(handle.getvalue()) == [] + + +class TestDirectoryExportSink: + @pytest.fixture() + def source_file(self, tmp_path: Path) -> Path: + src: Path = tmp_path / "src" / "doc.pdf" + src.parent.mkdir(parents=True) + src.write_bytes(b"PDF-CONTENT") + return src + + def test_add_file_copies_to_relative_arcname( + self, + tmp_path: Path, + source_file: Path, + ) -> None: + target: Path = tmp_path / "out" + target.mkdir() + with DirectoryExportSink( + target, + compare_checksums=False, + compare_json=False, + delete=False, + ) as sink: + sink.add_file(source_file, "originals/doc.pdf") + assert (target / "originals" / "doc.pdf").read_bytes() == b"PDF-CONTENT" + + def test_add_json_writes_file(self, tmp_path: Path) -> None: + target: Path = tmp_path / "out" + target.mkdir() + with DirectoryExportSink( + target, + compare_checksums=False, + compare_json=False, + delete=False, + ) as sink: + sink.add_json({"version": "x"}, "metadata.json") + assert json.loads((target / "metadata.json").read_text()) == {"version": "x"} + + def test_stream_writes_manifest(self, tmp_path: Path) -> None: + target: Path = tmp_path / "out" + target.mkdir() + with DirectoryExportSink( + target, + compare_checksums=False, + compare_json=False, + delete=False, + ) as sink: + with sink.stream("manifest.json") as handle: + writer: StreamingManifestWriter = StreamingManifestWriter(handle) + writer.write_record({"pk": 1}) + writer.close() + assert json.loads((target / "manifest.json").read_text()) == [{"pk": 1}] + + def test_add_file_skips_when_size_and_mtime_match( + self, + tmp_path: Path, + source_file: Path, + ) -> None: + # Pre-existing target with identical size+mtime but DIFFERENT content: + # if add_file skips (no compare-checksums), the old content survives. + target: Path = tmp_path / "out" + target.mkdir() + existing: Path = target / "originals" / "doc.pdf" + existing.parent.mkdir(parents=True) + # Same byte length as the source but different content + matching mtime, + # so a size/mtime comparison treats it as unchanged and skips the copy. + existing.write_bytes(b"X" * len(b"PDF-CONTENT")) + stat = source_file.stat() + os.utime(existing, (stat.st_atime, stat.st_mtime)) + with DirectoryExportSink( + target, + compare_checksums=False, + compare_json=False, + delete=False, + ) as sink: + sink.add_file(source_file, "originals/doc.pdf", checksum="abc") + assert existing.read_bytes() == b"X" * len(b"PDF-CONTENT") # skipped + + def test_add_file_recopies_when_compare_checksums_differ( + self, + tmp_path: Path, + source_file: Path, + ) -> None: + target: Path = tmp_path / "out" + target.mkdir() + existing: Path = target / "originals" / "doc.pdf" + existing.parent.mkdir(parents=True) + existing.write_bytes(b"X" * len(b"PDF-CONTENT")) + stat = source_file.stat() + os.utime(existing, (stat.st_atime, stat.st_mtime)) + with DirectoryExportSink( + target, + compare_checksums=True, + compare_json=False, + delete=False, + ) as sink: + # wrong checksum forces recopy despite matching size/mtime + sink.add_file(source_file, "originals/doc.pdf", checksum="not-the-real-sum") + assert existing.read_bytes() == b"PDF-CONTENT" # recopied + + def test_delete_prunes_unwritten_snapshot_files( + self, + tmp_path: Path, + source_file: Path, + ) -> None: + target: Path = tmp_path / "out" + target.mkdir() + stale: Path = target / "stale.pdf" + stale.write_bytes(b"STALE") + with DirectoryExportSink( + target, + compare_checksums=False, + compare_json=False, + delete=True, + ) as sink: + sink.add_file(source_file, "originals/doc.pdf") + assert not stale.exists() + assert (target / "originals" / "doc.pdf").exists() + + def test_no_delete_keeps_unwritten_files( + self, + tmp_path: Path, + source_file: Path, + ) -> None: + target: Path = tmp_path / "out" + target.mkdir() + stale: Path = target / "stale.pdf" + stale.write_bytes(b"STALE") + with DirectoryExportSink( + target, + compare_checksums=False, + compare_json=False, + delete=False, + ) as sink: + sink.add_file(source_file, "originals/doc.pdf") + assert stale.exists()