Add for split and merge too

src/documents/bulk_edit.py

@@ -476,12 +476,14 @@ def merge(
     metadata_document_id: int | None = None,
     delete_originals: bool = False,
     archive_fallback: bool = False,
+    source_mode: SourceMode = SourceModeChoices.LATEST_VERSION,
     user: User | None = None,
 ) -> Literal["OK"]:
     logger.info(
         f"Attempting to merge {len(doc_ids)} documents into a single document.",
     )
-    qs = Document.objects.filter(id__in=doc_ids)
+    qs = Document.objects.select_related("root_document").filter(id__in=doc_ids)
+    docs_by_id = {doc.id: doc for doc in qs}
     affected_docs: list[int] = []
     import pikepdf
 
@@ -490,14 +492,20 @@ def merge(
     handoff_asn: int | None = None
     # use doc_ids to preserve order
     for doc_id in doc_ids:
-        doc = qs.get(id=doc_id)
+        doc = docs_by_id.get(doc_id)
+        if doc is None:
+            continue
+        _, source_doc = _resolve_root_and_source_doc(
+            doc,
+            source_mode=source_mode,
+        )
         try:
             doc_path = (
-                doc.archive_path
+                source_doc.archive_path
                 if archive_fallback
-                and doc.mime_type != "application/pdf"
-                and doc.has_archive_version
-                else doc.source_path
+                and source_doc.mime_type != "application/pdf"
+                and source_doc.has_archive_version
+                else source_doc.source_path
             )
             with pikepdf.open(str(doc_path)) as pdf:
                 version = max(version, pdf.pdf_version)
@@ -579,18 +587,23 @@ def split(
     pages: list[list[int]],
     *,
     delete_originals: bool = False,
+    source_mode: SourceMode = SourceModeChoices.LATEST_VERSION,
     user: User | None = None,
 ) -> Literal["OK"]:
     logger.info(
         f"Attempting to split document {doc_ids[0]} into {len(pages)} documents",
     )
-    doc = Document.objects.get(id=doc_ids[0])
+    doc = Document.objects.select_related("root_document").get(id=doc_ids[0])
+    _, source_doc = _resolve_root_and_source_doc(
+        doc,
+        source_mode=source_mode,
+    )
     import pikepdf
 
     consume_tasks = []
 
     try:
-        with pikepdf.open(doc.source_path) as pdf:
+        with pikepdf.open(source_doc.source_path) as pdf:
             for idx, split_doc in enumerate(pages):
                 dst: pikepdf.Pdf = pikepdf.new()
                 for page in split_doc:

src/documents/tests/test_bulk_edit.py

@@ -660,6 +660,33 @@ class TestPDFActions(DirectoriesMixin, TestCase):
 
         self.assertEqual(result, "OK")
 
+    @mock.patch("pikepdf.open")
+    @mock.patch("documents.tasks.consume_file.s")
+    def test_merge_uses_latest_version_source_for_root_selection(
+        self,
+        mock_consume_file,
+        mock_open_pdf,
+    ) -> None:
+        version_file = self.dirs.scratch_dir / "sample2_version_merge.pdf"
+        shutil.copy(self.doc2.source_path, version_file)
+        version = Document.objects.create(
+            checksum="B-v1",
+            title="B version 1",
+            root_document=self.doc2,
+            filename=version_file,
+            mime_type="application/pdf",
+        )
+        fake_pdf = mock.MagicMock()
+        fake_pdf.pdf_version = "1.7"
+        fake_pdf.pages = [mock.Mock()]
+        mock_open_pdf.return_value.__enter__.return_value = fake_pdf
+
+        result = bulk_edit.merge([self.doc2.id])
+
+        self.assertEqual(result, "OK")
+        mock_open_pdf.assert_called_once_with(str(version.source_path))
+        mock_consume_file.assert_not_called()
+
     @mock.patch("documents.bulk_edit.delete.si")
     @mock.patch("documents.tasks.consume_file.s")
     def test_merge_and_delete_originals(
@@ -868,6 +895,36 @@ class TestPDFActions(DirectoriesMixin, TestCase):
 
         self.assertEqual(result, "OK")
 
+    @mock.patch("documents.bulk_edit.group")
+    @mock.patch("pikepdf.open")
+    @mock.patch("documents.tasks.consume_file.s")
+    def test_split_uses_latest_version_source_for_root_selection(
+        self,
+        mock_consume_file,
+        mock_open_pdf,
+        mock_group,
+    ) -> None:
+        version_file = self.dirs.scratch_dir / "sample2_version_split.pdf"
+        shutil.copy(self.doc2.source_path, version_file)
+        version = Document.objects.create(
+            checksum="B-v1",
+            title="B version 1",
+            root_document=self.doc2,
+            filename=version_file,
+            mime_type="application/pdf",
+        )
+        fake_pdf = mock.MagicMock()
+        fake_pdf.pages = [mock.Mock(), mock.Mock()]
+        mock_open_pdf.return_value.__enter__.return_value = fake_pdf
+        mock_group.return_value.delay.return_value = None
+
+        result = bulk_edit.split([self.doc2.id], [[1], [2]])
+
+        self.assertEqual(result, "OK")
+        mock_open_pdf.assert_called_once_with(version.source_path)
+        mock_consume_file.assert_not_called()
+        mock_group.return_value.delay.assert_not_called()
+
     @mock.patch("documents.bulk_edit.delete.si")
     @mock.patch("documents.tasks.consume_file.s")
     @mock.patch("documents.bulk_edit.chord")