From 8d956b30eb5acc210d58fc9a9321b1c99c82481f Mon Sep 17 00:00:00 2001 From: shamoon <4887959+shamoon@users.noreply.github.com> Date: Sat, 6 Jun 2026 21:01:46 -0700 Subject: [PATCH] Yea this makes more sense --- src/documents/bulk_edit.py | 120 +++++++++++++------------- src/documents/tests/test_bulk_edit.py | 66 +++++++++++--- 2 files changed, 114 insertions(+), 72 deletions(-) diff --git a/src/documents/bulk_edit.py b/src/documents/bulk_edit.py index ad6088cd1..0cea9a3a1 100644 --- a/src/documents/bulk_edit.py +++ b/src/documents/bulk_edit.py @@ -2,7 +2,6 @@ from __future__ import annotations import logging import tempfile -import warnings from pathlib import Path from typing import TYPE_CHECKING from typing import Literal @@ -905,13 +904,8 @@ def remove_password( doc.id, pair.source_doc.source_path, ) - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - message="A password was provided, but no password was needed to open this PDF.", - category=UserWarning, - ) - with pikepdf.open(source_path, password=password) as pdf: + try: + with pikepdf.open(source_path) as pdf: if not pdf.is_encrypted: logger.info( "Skipping password removal for document %s because the " @@ -919,63 +913,67 @@ def remove_password( pair.root_doc.id, ) continue + except pikepdf.PasswordError: + # Password-protected PDFs need the supplied password below. + pass - filepath: Path = ( - Path(tempfile.mkdtemp(dir=settings.SCRATCH_DIR)) - / f"{pair.root_doc.id}_unprotected.pdf" + with pikepdf.open(source_path, password=password) as pdf: + filepath: Path = ( + Path(tempfile.mkdtemp(dir=settings.SCRATCH_DIR)) + / f"{pair.root_doc.id}_unprotected.pdf" + ) + pdf.remove_unreferenced_resources() + pdf.save(filepath) + + if update_document: + # Create a new version rather than modifying the root/original in place. + overrides = ( + DocumentMetadataOverrides().from_document(pair.root_doc) + if include_metadata + else DocumentMetadataOverrides() ) - pdf.remove_unreferenced_resources() - pdf.save(filepath) + if user is not None: + overrides.owner_id = user.id + overrides.actor_id = user.id + consume_file.apply_async( + kwargs={ + "input_doc": ConsumableDocument( + source=DocumentSource.ConsumeFolder, + original_file=filepath, + root_document_id=pair.root_doc.id, + ), + "overrides": overrides, + }, + headers={"trigger_source": trigger_source}, + ) + else: + consume_tasks = [] + overrides = ( + DocumentMetadataOverrides().from_document(pair.root_doc) + if include_metadata + else DocumentMetadataOverrides() + ) + if user is not None: + overrides.owner_id = user.id + overrides.actor_id = user.id - if update_document: - # Create a new version rather than modifying the root/original in place. - overrides = ( - DocumentMetadataOverrides().from_document(pair.root_doc) - if include_metadata - else DocumentMetadataOverrides() - ) - if user is not None: - overrides.owner_id = user.id - overrides.actor_id = user.id - consume_file.apply_async( - kwargs={ - "input_doc": ConsumableDocument( - source=DocumentSource.ConsumeFolder, - original_file=filepath, - root_document_id=pair.root_doc.id, - ), - "overrides": overrides, - }, - headers={"trigger_source": trigger_source}, - ) + consume_tasks.append( + consume_file.s( + input_doc=ConsumableDocument( + source=DocumentSource.ConsumeFolder, + original_file=filepath, + ), + overrides=overrides, + ).set(headers={"trigger_source": trigger_source}), + ) + + if delete_original: + chord( + header=consume_tasks, + body=delete.si([doc.id]), + ).delay() else: - consume_tasks = [] - overrides = ( - DocumentMetadataOverrides().from_document(pair.root_doc) - if include_metadata - else DocumentMetadataOverrides() - ) - if user is not None: - overrides.owner_id = user.id - overrides.actor_id = user.id - - consume_tasks.append( - consume_file.s( - input_doc=ConsumableDocument( - source=DocumentSource.ConsumeFolder, - original_file=filepath, - ), - overrides=overrides, - ).set(headers={"trigger_source": trigger_source}), - ) - - if delete_original: - chord( - header=consume_tasks, - body=delete.si([doc.id]), - ).delay() - else: - group(consume_tasks).delay() + group(consume_tasks).delay() except Exception as e: logger.exception( diff --git a/src/documents/tests/test_bulk_edit.py b/src/documents/tests/test_bulk_edit.py index 27c874579..010744af1 100644 --- a/src/documents/tests/test_bulk_edit.py +++ b/src/documents/tests/test_bulk_edit.py @@ -3,6 +3,7 @@ from datetime import date from pathlib import Path from unittest import mock +import pikepdf from django.contrib.auth.models import Group from django.contrib.auth.models import User from django.test import TestCase @@ -615,6 +616,18 @@ class TestPDFActions(DirectoriesMixin, TestCase): self.img_doc.archive_filename = img_doc_archive self.img_doc.save() + @staticmethod + def mock_password_required_pdf( + mock_open: mock.Mock, + fake_pdf: mock.Mock, + ) -> None: + password_context = mock.MagicMock() + password_context.__enter__.return_value = fake_pdf + mock_open.side_effect = [ + pikepdf.PasswordError("password required"), + password_context, + ] + @mock.patch("documents.tasks.consume_file.s") def test_merge(self, mock_consume_file) -> None: """ @@ -1466,6 +1479,7 @@ class TestPDFActions(DirectoriesMixin, TestCase): fake_pdf = mock.MagicMock() fake_pdf.pages = [mock.Mock(), mock.Mock(), mock.Mock()] + fake_pdf.is_encrypted = True def save_side_effect(target_path): Path(target_path).write_bytes(b"new pdf content") @@ -1480,7 +1494,13 @@ class TestPDFActions(DirectoriesMixin, TestCase): ) self.assertEqual(result, "OK") - mock_open.assert_called_once_with(doc.source_path, password="secret") + self.assertEqual( + mock_open.call_args_list, + [ + mock.call(doc.source_path), + mock.call(doc.source_path, password="secret"), + ], + ) fake_pdf.remove_unreferenced_resources.assert_called_once() mock_update_document.assert_not_called() mock_consume_delay.assert_called_once() @@ -1515,7 +1535,7 @@ class TestPDFActions(DirectoriesMixin, TestCase): ) self.assertEqual(result, "OK") - mock_open.assert_called_once_with(doc.source_path, password="secret") + mock_open.assert_called_once_with(doc.source_path) fake_pdf.remove_unreferenced_resources.assert_not_called() fake_pdf.save.assert_not_called() mock_mkdtemp.assert_not_called() @@ -1540,12 +1560,12 @@ class TestPDFActions(DirectoriesMixin, TestCase): mock_mkdtemp.return_value = str(temp_dir) fake_pdf = mock.MagicMock() + self.mock_password_required_pdf(mock_open, fake_pdf) def save_side_effect(target_path): Path(target_path).write_bytes(b"new pdf content") fake_pdf.save.side_effect = save_side_effect - mock_open.return_value.__enter__.return_value = fake_pdf result = bulk_edit.remove_password( [doc.id], @@ -1555,7 +1575,13 @@ class TestPDFActions(DirectoriesMixin, TestCase): ) self.assertEqual(result, "OK") - mock_open.assert_called_once_with(source_file, password="secret") + self.assertEqual( + mock_open.call_args_list, + [ + mock.call(source_file), + mock.call(source_file, password="secret"), + ], + ) mock_update_document.assert_not_called() mock_consume_delay.assert_called_once() @@ -1574,7 +1600,7 @@ class TestPDFActions(DirectoriesMixin, TestCase): root_document=self.doc1, ) fake_pdf = mock.MagicMock() - mock_open.return_value.__enter__.return_value = fake_pdf + self.mock_password_required_pdf(mock_open, fake_pdf) result = bulk_edit.remove_password( [self.doc1.id], @@ -1584,7 +1610,13 @@ class TestPDFActions(DirectoriesMixin, TestCase): ) self.assertEqual(result, "OK") - mock_open.assert_called_once_with(self.doc1.source_path, password="secret") + self.assertEqual( + mock_open.call_args_list, + [ + mock.call(self.doc1.source_path), + mock.call(self.doc1.source_path, password="secret"), + ], + ) mock_consume_delay.assert_called_once() @mock.patch("documents.bulk_edit.chord") @@ -1607,12 +1639,12 @@ class TestPDFActions(DirectoriesMixin, TestCase): fake_pdf = mock.MagicMock() fake_pdf.pages = [mock.Mock(), mock.Mock()] + self.mock_password_required_pdf(mock_open, fake_pdf) def save_side_effect(target_path: Path) -> None: target_path.write_bytes(b"password removed") fake_pdf.save.side_effect = save_side_effect - mock_open.return_value.__enter__.return_value = fake_pdf mock_group.return_value.delay.return_value = None user = User.objects.create(username="owner") @@ -1627,7 +1659,13 @@ class TestPDFActions(DirectoriesMixin, TestCase): ) self.assertEqual(result, "OK") - mock_open.assert_called_once_with(doc.source_path, password="secret") + self.assertEqual( + mock_open.call_args_list, + [ + mock.call(doc.source_path), + mock.call(doc.source_path, password="secret"), + ], + ) mock_consume_file.assert_called_once() call_kwargs = mock_consume_file.call_args.kwargs consumable_document = call_kwargs["input_doc"] @@ -1673,7 +1711,7 @@ class TestPDFActions(DirectoriesMixin, TestCase): ) self.assertEqual(result, "OK") - mock_open.assert_called_once_with(doc.source_path, password="secret") + mock_open.assert_called_once_with(doc.source_path) fake_pdf.remove_unreferenced_resources.assert_not_called() fake_pdf.save.assert_not_called() mock_mkdtemp.assert_not_called() @@ -1704,12 +1742,12 @@ class TestPDFActions(DirectoriesMixin, TestCase): fake_pdf = mock.MagicMock() fake_pdf.pages = [mock.Mock(), mock.Mock()] + self.mock_password_required_pdf(mock_open, fake_pdf) def save_side_effect(target_path: Path) -> None: target_path.write_bytes(b"password removed") fake_pdf.save.side_effect = save_side_effect - mock_open.return_value.__enter__.return_value = fake_pdf mock_chord.return_value.delay.return_value = None result = bulk_edit.remove_password( @@ -1721,7 +1759,13 @@ class TestPDFActions(DirectoriesMixin, TestCase): ) self.assertEqual(result, "OK") - mock_open.assert_called_once_with(doc.source_path, password="secret") + self.assertEqual( + mock_open.call_args_list, + [ + mock.call(doc.source_path), + mock.call(doc.source_path, password="secret"), + ], + ) mock_consume_file.assert_called_once() mock_group.assert_not_called() mock_chord.assert_called_once()