From bbdc016a328a10d52dc14760f0a70e700220d90b Mon Sep 17 00:00:00 2001 From: shamoon <4887959+shamoon@users.noreply.github.com> Date: Tue, 20 Oct 2026 13:51:54 -0700 Subject: [PATCH] Backend for manual tempalte choice --- src/documents/bulk_edit.py | 14 ++-- src/documents/serialisers.py | 10 ++- src/documents/tasks.py | 26 ++++++- src/documents/tests/test_api_bulk_edit.py | 60 +++++++++++++++++ src/documents/tests/test_api_ocr_templates.py | 67 +++++++++++++++++++ src/documents/tests/test_bulk_edit.py | 66 +++++++++++++++++- src/documents/views.py | 31 ++++++++- src/documents/zone_ocr.py | 8 ++- 8 files changed, 265 insertions(+), 17 deletions(-) diff --git a/src/documents/bulk_edit.py b/src/documents/bulk_edit.py index 42aeef891..24bc2cf7b 100644 --- a/src/documents/bulk_edit.py +++ b/src/documents/bulk_edit.py @@ -31,6 +31,7 @@ from documents.models import CustomField from documents.models import CustomFieldInstance from documents.models import Document from documents.models import DocumentType +from documents.models import OcrTemplate from documents.models import PaperlessTask from documents.models import StoragePath from documents.models import Tag @@ -1234,14 +1235,19 @@ def remove_doclink( Document.objects.filter(id=target_doc_id).update(modified=timezone.now()) -def run_zone_ocr(doc_ids: list[int]) -> Literal["OK"]: +def run_zone_ocr( + doc_ids: list[int], + *, + template: OcrTemplate | None = None, +) -> Literal["OK"]: """ - Queue zone OCR for the given documents. Templates are matched against each - document's content when its task runs. + Queue zone OCR for the given documents, with the given template or with the + templates that match each document's content when its task runs. """ + template_id = template.pk if template is not None else None for document_id in doc_ids: extract_zones.apply_async( - kwargs={"document_id": document_id}, + kwargs={"document_id": document_id, "template_id": template_id}, headers={"trigger_source": PaperlessTask.TriggerSource.MANUAL}, ) diff --git a/src/documents/serialisers.py b/src/documents/serialisers.py index 73963b757..77981da25 100644 --- a/src/documents/serialisers.py +++ b/src/documents/serialisers.py @@ -1856,7 +1856,15 @@ class ReprocessDocumentsSerializer(DocumentSelectionSerializer): remote_ocr = serializers.BooleanField(required=False, default=False) -class RunZoneOcrDocumentsSerializer(DocumentSelectionSerializer): +class RunZoneOcrSerializer(serializers.Serializer): + template = serializers.PrimaryKeyRelatedField( + queryset=OcrTemplate.objects.prefetch_related("zones", "zones__custom_field"), + required=False, + allow_null=True, + ) + + +class RunZoneOcrDocumentsSerializer(DocumentSelectionSerializer, RunZoneOcrSerializer): pass diff --git a/src/documents/tasks.py b/src/documents/tasks.py index 827c8f78b..e7997b02c 100644 --- a/src/documents/tasks.py +++ b/src/documents/tasks.py @@ -47,6 +47,7 @@ from documents.models import CustomFieldInstance from documents.models import Document from documents.models import DocumentBarcode from documents.models import DocumentType +from documents.models import OcrTemplate from documents.models import PaperlessTask from documents.models import ShareLink from documents.models import ShareLinkBundle @@ -370,9 +371,28 @@ def _read_barcodes_for_reprocess(document: Document) -> list[StoredBarcode] | No @shared_task -def extract_zones(document_id: int) -> None: - """Run the document type's enabled OCR templates against a stored document.""" - run_zone_extraction(Document.objects.get(id=document_id), None) +def extract_zones(document_id: int, template_id: int | None = None) -> None: + """ + Run an OCR template against a stored document, or every template that + matches it when no template is given. + """ + document = Document.objects.get(id=document_id) + if template_id is None: + run_zone_extraction(document, None) + return + + try: + template = OcrTemplate.objects.prefetch_related( + "zones", + "zones__custom_field", + ).get(id=template_id) + except OcrTemplate.DoesNotExist: + logger.warning( + "OCR template %s no longer exists, not running zone OCR", + template_id, + ) + return + run_zone_extraction(document, None, templates=[template]) @shared_task diff --git a/src/documents/tests/test_api_bulk_edit.py b/src/documents/tests/test_api_bulk_edit.py index 0bfb6f7c2..68c2a5a8f 100644 --- a/src/documents/tests/test_api_bulk_edit.py +++ b/src/documents/tests/test_api_bulk_edit.py @@ -12,6 +12,7 @@ from documents.models import CustomField from documents.models import CustomFieldInstance from documents.models import Document from documents.models import DocumentType +from documents.models import OcrTemplate from documents.models import StoragePath from documents.models import Tag from paperless_testing.dirs import DirectoriesMixin @@ -605,6 +606,65 @@ class TestBulkEditAPI(DirectoriesMixin, APITestCase): self.assertEqual(args[0], [self.doc1.id]) self.assertEqual(kwargs, {}) + @mock.patch("documents.views.bulk_edit.run_zone_ocr") + def test_run_zone_ocr_documents_endpoint_with_template(self, m) -> None: + """ + GIVEN: + - API data to run a chosen OCR template on a document + WHEN: + - API is called + THEN: + - run_zone_ocr is called with the template + """ + self.setup_mock(m, "run_zone_ocr") + template = OcrTemplate.objects.create( + name="Bills", + source_width=100, + source_height=100, + ) + response = self.client.post( + "/api/documents/run_zone_ocr/", + json.dumps({"documents": [self.doc1.id], "template": template.pk}), + content_type="application/json", + ) + self.assertEqual(response.status_code, status.HTTP_200_OK) + m.assert_called_once() + args, kwargs = m.call_args + self.assertEqual(args[0], [self.doc1.id]) + self.assertEqual(kwargs, {"template": template}) + + @mock.patch("documents.views.bulk_edit.run_zone_ocr") + def test_run_zone_ocr_documents_endpoint_rejects_private_template( + self, + m, + ) -> None: + """ + GIVEN: + - A user who can change the documents but not view the template + WHEN: + - API is called with that template + THEN: + - The whole request is refused + """ + user = UserFactory(username="editor") + grant_global(user, "view_document", "change_document") + template = OcrTemplate.objects.create( + name="Private", + owner=UserFactory(username="other"), + source_width=100, + source_height=100, + ) + self.client.force_authenticate(user=user) + + response = self.client.post( + "/api/documents/run_zone_ocr/", + json.dumps({"documents": [self.doc1.id], "template": template.pk}), + content_type="application/json", + ) + + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + m.assert_not_called() + @mock.patch("documents.serialisers.bulk_edit.set_storage_path") def test_api_set_storage_path(self, m) -> None: """ diff --git a/src/documents/tests/test_api_ocr_templates.py b/src/documents/tests/test_api_ocr_templates.py index b5def1496..5c4f90c97 100644 --- a/src/documents/tests/test_api_ocr_templates.py +++ b/src/documents/tests/test_api_ocr_templates.py @@ -561,6 +561,73 @@ class TestOcrTemplatesAPI(DirectoriesMixin, APITestCase): ) extract.assert_not_called() + @patch("documents.zone_ocr.run_zone_extraction") + def test_run_zone_ocr_applies_a_chosen_template(self, extract): + """ + GIVEN: + - A disabled template that does not match the document + WHEN: + - Zone OCR is run on the document with that template chosen + THEN: + - The template is applied anyway, since it was chosen explicitly + """ + document = DocumentFactory(content=CONTENT, mime_type="application/pdf") + document.source_path.write_bytes(b"test PDF") + template = OcrTemplate.objects.create( + name="Manual only", + match="Contoso", + enabled=False, + source_width=1000, + source_height=1000, + ) + + response = self.client.post( + f"/api/documents/{document.pk}/run-zone-ocr/", + {"template": template.pk}, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual(response.json()["status"], "completed") + extract.assert_called_once_with(document, None, templates=[template]) + + @patch("documents.zone_ocr.run_zone_extraction") + def test_run_zone_ocr_rejects_a_template_the_user_cannot_view(self, extract): + user = User.objects.create_user(username="editor") + user.user_permissions.add( + *Permission.objects.filter( + codename__in=["view_document", "change_document"], + ), + ) + document = DocumentFactory(content=CONTENT, owner=user) + template = OcrTemplate.objects.create( + name="Private", + owner=User.objects.create_user(username="other"), + source_width=1000, + source_height=1000, + ) + self.client.force_authenticate(user=user) + + response = self.client.post( + f"/api/documents/{document.pk}/run-zone-ocr/", + {"template": template.pk}, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + extract.assert_not_called() + + def test_run_zone_ocr_rejects_an_unknown_template(self): + document = DocumentFactory(content=CONTENT) + + response = self.client.post( + f"/api/documents/{document.pk}/run-zone-ocr/", + {"template": 999}, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + def test_document_page_image_requires_view_permission_on_the_document(self): owner = User.objects.create_user(username="owner") user = User.objects.create_user(username="template_editor") diff --git a/src/documents/tests/test_bulk_edit.py b/src/documents/tests/test_bulk_edit.py index bd569a965..09d2d6111 100644 --- a/src/documents/tests/test_bulk_edit.py +++ b/src/documents/tests/test_bulk_edit.py @@ -21,6 +21,7 @@ from documents.models import CustomField from documents.models import CustomFieldInstance from documents.models import Document from documents.models import DocumentType +from documents.models import OcrTemplate from documents.models import StoragePath from documents.models import Tag from documents.permissions import set_permissions_for_objects @@ -1999,9 +2000,31 @@ class TestBulkEditRunZoneOcr: assert result == "OK" assert [ - call.kwargs["kwargs"]["document_id"] - for call in mock_task.apply_async.call_args_list - ] == [doc.id for doc in docs] + call.kwargs["kwargs"] for call in mock_task.apply_async.call_args_list + ] == [{"document_id": doc.id, "template_id": None} for doc in docs] + + def test_queues_the_chosen_template(self, mock_task: mock.MagicMock) -> None: + """ + GIVEN: + - A chosen OCR template + WHEN: + - run_zone_ocr is called with it + THEN: + - The template is queued with the document + """ + doc = DocumentFactory() + template = OcrTemplate.objects.create( + name="Bills", + source_width=100, + source_height=100, + ) + + bulk_edit.run_zone_ocr([doc.id], template=template) + + assert mock_task.apply_async.call_args.kwargs["kwargs"] == { + "document_id": doc.id, + "template_id": template.pk, + } def test_task_runs_extraction(self, mocker: MockerFixture) -> None: """ @@ -2019,6 +2042,43 @@ class TestBulkEditRunZoneOcr: extract.assert_called_once_with(doc, None) + def test_task_runs_the_chosen_template(self, mocker: MockerFixture) -> None: + """ + GIVEN: + - A queued zone OCR task with a chosen template + WHEN: + - The task runs + THEN: + - Zone extraction runs with only that template + """ + extract = mocker.patch("documents.tasks.run_zone_extraction") + doc = DocumentFactory() + template = OcrTemplate.objects.create( + name="Bills", + source_width=100, + source_height=100, + ) + + extract_zones(doc.id, template.pk) + + extract.assert_called_once_with(doc, None, templates=[template]) + + def test_task_skips_a_deleted_template(self, mocker: MockerFixture) -> None: + """ + GIVEN: + - A template deleted between queueing and running + WHEN: + - The task runs + THEN: + - Nothing is extracted + """ + extract = mocker.patch("documents.tasks.run_zone_extraction") + doc = DocumentFactory() + + extract_zones(doc.id, 999) + + extract.assert_not_called() + @pytest.mark.django_db class TestBulkEditReprocess: diff --git a/src/documents/views.py b/src/documents/views.py index 9b0cc82bd..801ec141c 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -214,6 +214,7 @@ from documents.serialisers import ReprocessDocumentsSerializer from documents.serialisers import RotateDocumentsSerializer from documents.serialisers import RunTaskSerializer from documents.serialisers import RunZoneOcrDocumentsSerializer +from documents.serialisers import RunZoneOcrSerializer from documents.serialisers import SavedViewSerializer from documents.serialisers import SearchResultSerializer from documents.serialisers import SerializerWithPerms @@ -2281,6 +2282,7 @@ class DocumentViewSet( @extend_schema( operation_id="documents_run_zone_ocr", + request=RunZoneOcrSerializer, responses=inline_serializer( name="RunZoneOcrResult", fields={ @@ -2295,7 +2297,14 @@ class DocumentViewSet( ) @action(methods=["post"], detail=True, url_path="run-zone-ocr") def run_zone_ocr(self, request, pk=None): - """Run zone-based OCR extraction on this document.""" + """ + Run zone-based OCR extraction on this document, with the given template + or with every template that matches it. + """ + serializer = RunZoneOcrSerializer(data=request.data) + serializer.is_valid(raise_exception=True) + template = serializer.validated_data.get("template") + try: document = Document.objects.select_related( "root_document__owner", @@ -2307,9 +2316,18 @@ class DocumentViewSet( ) or not has_perms_owner_aware(request.user, "change_document", document): return HttpResponseForbidden("Insufficient permissions") + if template is not None and not has_perms_owner_aware( + request.user, + "view_ocrtemplate", + template, + ): + return HttpResponseForbidden("Insufficient permissions") + from documents.zone_ocr import get_templates_for_document - templates = get_templates_for_document(document) + templates = ( + [template] if template is not None else get_templates_for_document(document) + ) if not templates: return Response( {"status": "skipped", "reason": "no_template", "results": []}, @@ -2324,7 +2342,7 @@ class DocumentViewSet( from documents.zone_ocr import run_zone_extraction - run_zone_extraction(document, None) + run_zone_extraction(document, None, templates=templates) # Collect results results = [] @@ -3530,6 +3548,13 @@ class RunZoneOcrDocumentsView(DocumentOperationPermissionMixin): def post(self, request, *args, **kwargs): serializer = self.get_serializer(data=request.data) serializer.is_valid(raise_exception=True) + template = serializer.validated_data.get("template") + if template is not None and not has_perms_owner_aware( + request.user, + "view_ocrtemplate", + template, + ): + return HttpResponseForbidden("Insufficient permissions") return self._execute_document_action( method=bulk_edit.run_zone_ocr, validated_data=serializer.validated_data, diff --git a/src/documents/zone_ocr.py b/src/documents/zone_ocr.py index a3594cc77..c99ad0ffa 100644 --- a/src/documents/zone_ocr.py +++ b/src/documents/zone_ocr.py @@ -51,12 +51,14 @@ def get_templates_for_document(document: Document) -> list[OcrTemplate]: def run_zone_extraction( document: Document, original_file: Path | None, + templates: list[OcrTemplate] | None = None, ) -> None: """ - Run zone-based OCR extraction for every template that matches a document. - Called from the document_consumption_finished signal handler. + Run zone-based OCR extraction for the given templates, or for every + template that matches the document when none are given. """ - templates = get_templates_for_document(document) + if templates is None: + templates = get_templates_for_document(document) if not templates: return