diff --git a/src/documents/bulk_edit.py b/src/documents/bulk_edit.py index ec0d77637..7d110ab72 100644 --- a/src/documents/bulk_edit.py +++ b/src/documents/bulk_edit.py @@ -38,6 +38,7 @@ from documents.permissions import set_permissions_for_objects from documents.plugins.helpers import DocumentsStatusManager from documents.tasks import bulk_update_documents from documents.tasks import consume_file +from documents.tasks import extract_zones from documents.tasks import remove_document_from_index from documents.tasks import update_document_content_maybe_archive_file from documents.versioning import get_latest_version_for_root @@ -1231,3 +1232,26 @@ def remove_doclink( target_doc_field_instance.value.remove(document.id) target_doc_field_instance.save() Document.objects.filter(id=target_doc_id).update(modified=timezone.now()) + + +def run_zone_ocr(doc_ids: list[int]) -> Literal["OK"]: + """ + Queue zone OCR for the given documents. Documents whose type has no enabled + OCR template are skipped. + """ + document_ids = ( + Document.objects.filter( + id__in=doc_ids, + document_type__ocr_templates__enabled=True, + ) + .order_by() + .values_list("id", flat=True) + .distinct() + ) + for document_id in document_ids: + extract_zones.apply_async( + kwargs={"document_id": document_id}, + headers={"trigger_source": PaperlessTask.TriggerSource.MANUAL}, + ) + + return "OK" diff --git a/src/documents/serialisers.py b/src/documents/serialisers.py index 658b7253c..70ca2d9b1 100644 --- a/src/documents/serialisers.py +++ b/src/documents/serialisers.py @@ -1856,6 +1856,10 @@ class ReprocessDocumentsSerializer(DocumentSelectionSerializer): remote_ocr = serializers.BooleanField(required=False, default=False) +class RunZoneOcrDocumentsSerializer(DocumentSelectionSerializer): + pass + + class BulkEditSerializer( SerializerWithPerms, DocumentSelectionSerializer, diff --git a/src/documents/tasks.py b/src/documents/tasks.py index 778ee9673..827c8f78b 100644 --- a/src/documents/tasks.py +++ b/src/documents/tasks.py @@ -69,6 +69,7 @@ from documents.utils import compute_checksum from documents.utils import identity from documents.versioning import annotate_effective_content from documents.workflows.utils import get_workflows_for_trigger +from documents.zone_ocr import run_zone_extraction from paperless.config import AIConfig from paperless.config import BarcodeConfig from paperless.config import RemoteOCRConfig @@ -368,6 +369,12 @@ def _read_barcodes_for_reprocess(document: Document) -> list[StoredBarcode] | No return None +@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) + + @shared_task def update_document_content_maybe_archive_file( document_id, diff --git a/src/documents/tests/test_api_bulk_edit.py b/src/documents/tests/test_api_bulk_edit.py index 208657564..0bfb6f7c2 100644 --- a/src/documents/tests/test_api_bulk_edit.py +++ b/src/documents/tests/test_api_bulk_edit.py @@ -583,6 +583,28 @@ class TestBulkEditAPI(DirectoriesMixin, APITestCase): self.assertEqual(args[0], [self.doc1.id]) self.assertEqual(kwargs, {"remote_ocr": True}) + @mock.patch("documents.views.bulk_edit.run_zone_ocr") + def test_run_zone_ocr_documents_endpoint(self, m) -> None: + """ + GIVEN: + - API data to run zone OCR on a document + WHEN: + - API is called + THEN: + - run_zone_ocr is called with the document + """ + self.setup_mock(m, "run_zone_ocr") + response = self.client.post( + "/api/documents/run_zone_ocr/", + json.dumps({"documents": [self.doc1.id]}), + 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, {}) + @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 c474571bb..4fa9ba445 100644 --- a/src/documents/tests/test_api_ocr_templates.py +++ b/src/documents/tests/test_api_ocr_templates.py @@ -3,6 +3,7 @@ import json from unittest.mock import patch +from django.contrib.auth.models import Permission from django.contrib.auth.models import User from rest_framework import status from rest_framework.test import APITestCase @@ -437,6 +438,26 @@ class TestOcrTemplatesAPI(DirectoriesMixin, APITestCase): ) extract.assert_not_called() + @patch("documents.zone_ocr.run_zone_extraction") + def test_run_zone_ocr_requires_change_permission_on_the_document( + self, + extract, + ): + owner = User.objects.create_user(username="owner") + user = User.objects.create_user(username="editor") + user.user_permissions.add( + *Permission.objects.filter( + codename__in=["view_document", "change_document"], + ), + ) + document = DocumentFactory(document_type=self.doc_type, owner=owner) + self.client.force_authenticate(user=user) + + response = self.client.post(f"/api/documents/{document.pk}/run-zone-ocr/") + + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + extract.assert_not_called() + def test_run_zone_ocr_missing_document_is_an_error(self): response = self.client.post("/api/documents/99999/run-zone-ocr/") diff --git a/src/documents/tests/test_bulk_edit.py b/src/documents/tests/test_bulk_edit.py index 5aa760032..ec8759dbd 100644 --- a/src/documents/tests/test_bulk_edit.py +++ b/src/documents/tests/test_bulk_edit.py @@ -21,9 +21,11 @@ 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 +from documents.tasks import extract_zones from paperless_testing.dirs import DirectoriesMixin from paperless_testing.factories import DocumentFactory from paperless_testing.permissions import grant_object @@ -1973,6 +1975,69 @@ class TestPDFActions(DirectoriesMixin, TestCase): self.assertIn("Error removing password from document", cm.output[0]) +@pytest.mark.django_db +class TestBulkEditRunZoneOcr: + @pytest.fixture + def mock_task(self, mocker: MockerFixture) -> mock.MagicMock: + return mocker.patch("documents.bulk_edit.extract_zones") + + def test_queues_documents_with_an_enabled_template( + self, + mock_task: mock.MagicMock, + ) -> None: + """ + GIVEN: + - Documents whose type has enabled templates, a disabled template, + or no type at all + WHEN: + - run_zone_ocr is called with all of them + THEN: + - Only the document with an enabled template is queued, once + """ + invoice = DocumentType.objects.create(name="Invoice") + receipt = DocumentType.objects.create(name="Receipt") + for name in ("Invoices", "More invoices"): + OcrTemplate.objects.create( + name=name, + document_type=invoice, + source_width=1000, + source_height=1000, + ) + OcrTemplate.objects.create( + name="Receipts", + document_type=receipt, + source_width=1000, + source_height=1000, + enabled=False, + ) + templated = DocumentFactory(document_type=invoice) + others = [DocumentFactory(document_type=receipt), DocumentFactory()] + + result = bulk_edit.run_zone_ocr([templated.id, *(d.id for d in others)]) + + assert result == "OK" + mock_task.apply_async.assert_called_once() + assert mock_task.apply_async.call_args.kwargs["kwargs"] == { + "document_id": templated.id, + } + + def test_task_runs_extraction(self, mocker: MockerFixture) -> None: + """ + GIVEN: + - A queued zone OCR task + WHEN: + - The task runs + THEN: + - Zone extraction runs on the stored document + """ + extract = mocker.patch("documents.tasks.run_zone_extraction") + doc = DocumentFactory() + + extract_zones(doc.id) + + extract.assert_called_once_with(doc, None) + + @pytest.mark.django_db class TestBulkEditReprocess: @pytest.fixture diff --git a/src/documents/views.py b/src/documents/views.py index 622706ed7..427e99305 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -212,6 +212,7 @@ from documents.serialisers import RemovePasswordDocumentsSerializer from documents.serialisers import ReprocessDocumentsSerializer from documents.serialisers import RotateDocumentsSerializer from documents.serialisers import RunTaskSerializer +from documents.serialisers import RunZoneOcrDocumentsSerializer from documents.serialisers import SavedViewSerializer from documents.serialisers import SearchResultSerializer from documents.serialisers import SerializerWithPerms @@ -2298,6 +2299,8 @@ class DocumentViewSet( document = Document.objects.get(pk=pk) except Document.DoesNotExist: raise Http404 + if not has_perms_owner_aware(request.user, "change_document", document): + return HttpResponseForbidden("Insufficient permissions") if not document.document_type_id: return Response( @@ -3508,6 +3511,33 @@ class ReprocessDocumentsView(DocumentOperationPermissionMixin): ) +@extend_schema_view( + post=extend_schema( + operation_id="documents_run_zone_ocr_bulk", + description="Run OCR template zone extraction on selected documents", + responses={ + 200: inline_serializer( + name="RunZoneOcrDocumentsResult", + fields={ + "result": serializers.CharField(), + }, + ), + }, + ), +) +class RunZoneOcrDocumentsView(DocumentOperationPermissionMixin): + serializer_class = RunZoneOcrDocumentsSerializer + + def post(self, request, *args, **kwargs): + serializer = self.get_serializer(data=request.data) + serializer.is_valid(raise_exception=True) + return self._execute_document_action( + method=bulk_edit.run_zone_ocr, + validated_data=serializer.validated_data, + operation_label="document zone OCR", + ) + + @extend_schema_view( post=extend_schema( operation_id="documents_edit_pdf", diff --git a/src/paperless/urls.py b/src/paperless/urls.py index 6733ea0fa..b3dcdbdf9 100644 --- a/src/paperless/urls.py +++ b/src/paperless/urls.py @@ -36,6 +36,7 @@ from documents.views import RemoteVersionView from documents.views import RemovePasswordDocumentsView from documents.views import ReprocessDocumentsView from documents.views import RotateDocumentsView +from documents.views import RunZoneOcrDocumentsView from documents.views import SavedViewViewSet from documents.views import SearchAutoCompleteView from documents.views import SelectionDataView @@ -168,6 +169,11 @@ urlpatterns = [ ReprocessDocumentsView.as_view(), name="reprocess_documents", ), + re_path( + "^run_zone_ocr/", + RunZoneOcrDocumentsView.as_view(), + name="run_zone_ocr_documents", + ), re_path( "^rotate/", RotateDocumentsView.as_view(),