Backend for manual tempalte choice

This commit is contained in:
shamoon committed 2026-10-20 13:51:54 -07:00
1 parent eae392469f
commit bbdc016a32
8 files changed
+265 -17

No files matched your search

+10 -4
View File
@@ -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},
)
+9 -1
View File
@@ -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
+23 -3
View File
@@ -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
+60
View File
@@ -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:
"""
@@ -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")
+63 -3
View File
@@ -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:
+28 -3
View File
@@ -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,
+5 -3
View File
@@ -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