Backend work for proper bulk zone ocr

This commit is contained in:
shamoon committed 2026-10-20 10:20:28 -07:00
1 parent 7d9b8c770d
commit 8bb2f88015
8 files changed
+179

No files matched your search

+24
View File
@@ -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"
+4
View File
@@ -1856,6 +1856,10 @@ class ReprocessDocumentsSerializer(DocumentSelectionSerializer):
remote_ocr = serializers.BooleanField(required=False, default=False)
class RunZoneOcrDocumentsSerializer(DocumentSelectionSerializer):
pass
class BulkEditSerializer(
SerializerWithPerms,
DocumentSelectionSerializer,
+7
View File
@@ -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,
+22
View File
@@ -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:
"""
@@ -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/")
+65
View File
@@ -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
+30
View File
@@ -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",