diff --git a/src/documents/tests/test_zone_ocr.py b/src/documents/tests/test_zone_ocr.py index 81264d829..ed16da926 100644 --- a/src/documents/tests/test_zone_ocr.py +++ b/src/documents/tests/test_zone_ocr.py @@ -1,6 +1,7 @@ """Tests for the zone-based OCR extraction engine.""" import tempfile +from datetime import date from pathlib import Path from unittest.mock import MagicMock from unittest.mock import patch @@ -18,6 +19,7 @@ from documents.models import MatchingModel from documents.models import OcrTemplate from documents.models import OcrTemplateZone from documents.zone_ocr import _apply_transform +from documents.zone_ocr import _combine_field_value from documents.zone_ocr import _convert_value from documents.zone_ocr import _crop_zone from documents.zone_ocr import _detect_mime @@ -161,25 +163,150 @@ class TestZoneExtraction(SimpleTestCase): class TestSavedZoneExtraction(DirectoriesMixin, TestCase): - @patch("documents.zone_ocr.subprocess.run") - def test_saved_extraction_matches_preview_and_rejects_invalid_value(self, run): - page = Path(self.dirs.scratch_dir) / "page.png" + def setUp(self): + super().setUp() + self.page = self.dirs.scratch_dir / "page.png" with Image.new("RGB", (200, 100), "white") as image: - image.save(page) - document = DocumentFactory(title="Original") - template = OcrTemplate.objects.create( + image.save(self.page) + self.document = DocumentFactory(title="Original") + self.template = OcrTemplate.objects.create( name="Account", source_width=100, source_height=50, ) - zone = OcrTemplateZone.objects.create( - template=template, - name="Title", - target="title", + + def zone(self, name, target="title", **kwargs): + return OcrTemplateZone.objects.create( + template=self.template, + name=name, + target=target, x=0, y=0, width=20, height=10, + **kwargs, + ) + + @patch("documents.zone_ocr.subprocess.run") + def test_combines_zones_and_writes_each_field_once(self, run): + zones = [self.zone("First"), self.zone("Second")] + self.template.combine_formats = {"title": "{Second} / {First}"} + run.side_effect = [MagicMock(stdout="one"), MagicMock(stdout="two")] + with patch.object(self.document, "save", wraps=self.document.save) as save: + _process_template(self.document, self.page, self.template, zones) + self.document.refresh_from_db() + self.assertEqual(self.document.title, "two / one") + save.assert_called_once_with(update_fields=["title"]) + + @patch("documents.zone_ocr.subprocess.run") + def test_converts_custom_and_builtin_fields_before_saving(self, run): + field = CustomField.objects.create(name="Paid", data_type="boolean") + zones = [ + self.zone("Title"), + self.zone("ASN", "asn"), + self.zone("Created", "created"), + self.zone("Paid", "custom_field", custom_field=field), + ] + run.side_effect = [ + MagicMock(stdout=value) + for value in [ + "New title", + "ASN 42", + "2026-04-13", + "no", + ] + ] + _process_template(self.document, self.page, self.template, zones) + self.document.refresh_from_db() + self.assertEqual(self.document.title, "New title") + self.assertEqual(self.document.archive_serial_number, 42) + self.assertEqual(self.document.created, date(2026, 4, 13)) + self.assertFalse(self.document.custom_fields.get(field=field).value) + + @patch("documents.zone_ocr.subprocess.run") + def test_invalid_field_values_preserve_existing_values(self, run): + field = CustomField.objects.create(name="Number", data_type="integer") + CustomFieldInstance.objects.create( + document=self.document, + field=field, + value_int=12, + ) + zones = [ + self.zone("Title"), + self.zone("ASN", "asn"), + self.zone("Created", "created"), + self.zone("Number", "custom_field", custom_field=field), + self.zone("Unassigned", "custom_field"), + ] + original_created = self.document.created + run.side_effect = [ + MagicMock(stdout=value) + for value in [ + "Valid", + "9999999999999999999999999999999", + "invalid", + "abc", + "text", + ] + ] + _process_template(self.document, self.page, self.template, zones) + self.document.refresh_from_db() + self.assertEqual(self.document.title, "Valid") + self.assertIsNone(self.document.archive_serial_number) + self.assertEqual(self.document.created, original_created) + self.assertEqual(self.document.custom_fields.get(field=field).value, 12) + + @patch("documents.zone_ocr.subprocess.run") + def test_write_failure_rolls_back_template_and_restores_document(self, run): + field = CustomField.objects.create(name="Number", data_type="integer") + zones = [ + self.zone("Title"), + self.zone("Number", "custom_field", custom_field=field), + ] + run.side_effect = [MagicMock(stdout="New title"), MagicMock(stdout="42")] + save = self.document.save + + def fail_after_save(**kwargs): + save(**kwargs) + raise RuntimeError("save failed") + + with patch.object(self.document, "save", side_effect=fail_after_save): + with self.assertRaisesMessage(RuntimeError, "save failed"): + _process_template(self.document, self.page, self.template, zones) + self.assertEqual(self.document.title, "Original") + self.document.refresh_from_db() + self.assertEqual(self.document.title, "Original") + self.assertFalse(self.document.custom_fields.exists()) + + @patch("documents.zone_ocr.subprocess.run") + def test_last_successful_template_wins(self, run): + self.zone("First") + second = OcrTemplate.objects.create( + name="Second", + source_width=100, + source_height=50, + ) + self.template = second + self.zone("Second") + third = OcrTemplate.objects.create( + name="Third", + source_width=100, + source_height=50, + ) + self.template = third + self.zone("Third", validation_regex="accepted") + templates = list(OcrTemplate.objects.order_by("pk")) + run.side_effect = [ + MagicMock(stdout=value) for value in ["First", "Second", "rejected"] + ] + run_zone_extraction(self.document, self.page, templates=templates) + self.document.refresh_from_db() + self.assertEqual(self.document.title, "Second") + + @patch("documents.zone_ocr.subprocess.run") + def test_saved_extraction_matches_preview_and_rejects_invalid_value(self, run): + zone = self.zone( + "Title", transform="uppercase", validation_regex="ACCOUNT [A-Z]+", ) @@ -189,12 +316,30 @@ class TestSavedZoneExtraction(DirectoriesMixin, TestCase): ]: with self.subTest(raw_text=raw_text): run.return_value.stdout = raw_text - preview = extract_zone_preview(page, zone, 100, 50, 1) + preview = extract_zone_preview(self.page, zone, 100, 50, 1) self.assertEqual(preview["raw_text"], raw_text) self.assertEqual(preview["value"], raw_text.upper()) - _process_template(document, page, template, [zone]) - document.refresh_from_db() - self.assertEqual(document.title, expected_title) + _process_template(self.document, self.page, self.template, [zone]) + self.document.refresh_from_db() + self.assertEqual(self.document.title, expected_title) + + +class TestCombineFieldValue(SimpleTestCase): + def test_default_format_preserves_zone_order_and_duplicate_names(self): + zones = [ + (OcrTemplateZone(name="Same"), value) for value in ["one", None, "two"] + ] + self.assertEqual(_combine_field_value("", zones), "one two") + + def test_format_preserves_existing_missing_token_cleanup(self): + zones = [ + (OcrTemplateZone(name="First"), "one"), + (OcrTemplateZone(name="Second"), None), + ] + self.assertEqual( + _combine_field_value("{ First } / {Second} / {Unknown}", zones), + "one", + ) class TestTemplateVisibility(DirectoriesMixin, TestCase): diff --git a/src/documents/zone_ocr.py b/src/documents/zone_ocr.py index 9d659d490..1fbfebdde 100644 --- a/src/documents/zone_ocr.py +++ b/src/documents/zone_ocr.py @@ -20,6 +20,8 @@ from datetime import datetime from pathlib import Path from django.conf import settings +from django.core.exceptions import ValidationError +from django.db import transaction from PIL import Image from documents.matching import matches @@ -41,6 +43,13 @@ class ZoneExtraction: value: str | None = None +@dataclass(frozen=True) +class FieldUpdate: + field_name: str + value: object + custom_field: CustomField | None = None + + def get_templates_for_document(document: Document) -> list[OcrTemplate]: """ The enabled templates that match the document's content and that its owner @@ -140,14 +149,7 @@ def _process_template( template: OcrTemplate, zones: list[OcrTemplateZone], ) -> None: - """Process all zones in a template against a document. - - Each zone is OCR'd independently, then zones are grouped by their target - field and each field is written exactly once. When several zones share a - field, their values are combined via the template's per-field format string - (or joined in order if none is set) — this avoids the zones overwriting each - other's value. - """ + """Extract and combine zone values before applying a template's updates.""" pages_needed: set[int] = { _resolve_page_idx(zone.page, document.page_count) for zone in zones } @@ -162,30 +164,19 @@ def _process_template( document.page_count, ) - # Pass 1: OCR every zone into a value (or None if it failed/was rejected). - zone_values: dict[int, str | None] = {} + grouped: dict[str, list[tuple[OcrTemplateZone, str | None]]] = {} for zone in zones: page_idx = _resolve_page_idx(zone.page, document.page_count) - - if page_idx not in page_images: - logger.warning( - "Zone OCR: page %d not available for zone '%s'", - page_idx, - zone.name, + extracted = None + if page_idx in page_images: + result = _extract_zone( + page_images[page_idx], + zone, + zone.zone_source_width or template.source_width, + zone.zone_source_height or template.source_height, + tmp_path, ) - continue - - src_w = zone.zone_source_width or template.source_width - src_h = zone.zone_source_height or template.source_height - - result = _extract_zone( - page_images[page_idx], - zone, - src_w, - src_h, - tmp_path, - ) - extracted = result.value + extracted = result.value if ( extracted is not None @@ -200,38 +191,28 @@ def _process_template( ) extracted = None - zone_values[id(zone)] = extracted - - # Pass 2: group zones by target field and write each field once. - grouped: dict[str, list[OcrTemplateZone]] = {} - for zone in zones: - grouped.setdefault(_field_key(zone), []).append(zone) + grouped.setdefault(_field_key(zone), []).append((zone, extracted)) + updates = [] combine_formats = template.combine_formats or {} for key, field_zones in grouped.items(): value = _combine_field_value( combine_formats.get(key, ""), field_zones, - zone_values, ) if not value: continue - - target_zone = field_zones[0] - _write_zone_value(document, target_zone, value) - logger.info( - "Zone OCR: %s = %r (from %d zone(s))", - _zone_target_label(target_zone), - value[:100] if len(value) > 100 else value, - len(field_zones), - ) + update = _prepare_field_update(field_zones[0][0], value) + if update is not None: + updates.append(update) + _apply_field_updates(document, updates) def _field_key(zone: OcrTemplateZone) -> str: """Identify a zone's target field. Custom fields key by id, built-in targets by their name. Matches the key used in OcrTemplate.combine_formats and on the frontend field select.""" - target = getattr(zone, "target", None) or "custom_field" + target = zone.target if target == "custom_field" and zone.custom_field_id: return str(zone.custom_field_id) return target @@ -239,8 +220,7 @@ def _field_key(zone: OcrTemplateZone) -> str: def _combine_field_value( fmt: str, - field_zones: list[OcrTemplateZone], - zone_values: dict[int, str | None], + field_zones: list[tuple[OcrTemplateZone, str | None]], ) -> str: """Combine the OCR values of all zones targeting one field. @@ -248,11 +228,10 @@ def _combine_field_value( and literal text is kept; separators left dangling by an empty token are cleaned up. Without a format, the zone values are joined in order by a space. """ - values = {z.name: (zone_values.get(id(z)) or "") for z in field_zones} + values = {zone.name: value or "" for zone, value in field_zones} if not fmt: - parts = [zone_values.get(id(z)) or "" for z in field_zones] - return " ".join(p for p in parts if p).strip() + return " ".join(value for _, value in field_zones if value).strip() def _replace(match: re.Match) -> str: return values.get(match.group(1).strip(), "") @@ -531,14 +510,6 @@ def _apply_transform(text: str, transform: str, date_format: str = "") -> str: return text -def _zone_target_label(zone: OcrTemplateZone) -> str: - """Human label of a zone's write target (for logging).""" - target = getattr(zone, "target", None) or "custom_field" - if target == "custom_field": - return zone.custom_field.name if zone.custom_field_id else "(no field)" - return {"title": "Title", "asn": "ASN", "created": "Created"}.get(target, target) - - def _parse_created_datetime(value: str): """Parse an extracted value into a tz-aware datetime for document.created. @@ -568,70 +539,68 @@ def _parse_created_datetime(value: str): return None -def _write_zone_value( - document: Document, +def _prepare_field_update( zone: OcrTemplateZone, value: str, -) -> None: - """Write an extracted value to the zone's target — a custom field, or a - built-in document field (title / archive_serial_number / created).""" - target = getattr(zone, "target", None) or "custom_field" - - if target == "custom_field": - if zone.custom_field_id: - _write_custom_field(document, zone.custom_field, value) - else: - logger.debug("Zone OCR: zone '%s' has no custom field set", zone.name) - return - - if target == "title": - document.title = value[:128] - document.save(update_fields=["title"]) - elif target == "asn": +) -> FieldUpdate | None: + """Convert extracted text and enforce the destination field's constraints.""" + custom_field = None + if zone.target == OcrTemplateZone.TargetType.CUSTOM_FIELD: + if not zone.custom_field_id: + return None + custom_field = zone.custom_field + field_name = CustomFieldInstance.get_value_field_name(custom_field.data_type) + typed_value = _convert_value(value, custom_field.data_type) + model = CustomFieldInstance + elif zone.target == OcrTemplateZone.TargetType.TITLE: + field_name, typed_value = "title", value[:128] + model = Document + elif zone.target == OcrTemplateZone.TargetType.ASN: digits = re.sub(r"[^\d]", "", value) - if not digits: - logger.debug( - "Zone OCR: ASN zone '%s' produced no digits (%r)", - zone.name, - value[:50], - ) - return - document.archive_serial_number = int(digits) - document.save(update_fields=["archive_serial_number"]) - elif target == "created": + field_name = "archive_serial_number" + typed_value = int(digits) if digits else None + model = Document + elif zone.target == OcrTemplateZone.TargetType.CREATED: parsed = _parse_created_datetime(value) - if parsed is None: - logger.debug( - "Zone OCR: created zone '%s' could not parse a date (%r)", - zone.name, - value[:50], - ) - return - document.created = parsed - document.save(update_fields=["created"]) - - -def _write_custom_field( - document: Document, - custom_field: CustomField, - value: str, -) -> None: - """Write an extracted value to a document's custom field.""" - typed_value = _convert_value(value, custom_field.data_type) + field_name, typed_value = "created", parsed.date() if parsed else None + model = Document + else: + return None if typed_value is None: - logger.debug( - "Zone OCR: skipping custom field '%s' — value conversion returned None", - custom_field.name, - ) - return + return None + try: + typed_value = model._meta.get_field(field_name).clean(typed_value, None) + except ValidationError: + logger.info("Zone OCR: invalid value for zone '%s'", zone.name) + return None + return FieldUpdate(field_name, typed_value, custom_field) - value_field_name = CustomFieldInstance.get_value_field_name(custom_field.data_type) - CustomFieldInstance.objects.update_or_create( - document=document, - field=custom_field, - defaults={value_field_name: typed_value}, - ) +def _apply_field_updates(document: Document, updates: list[FieldUpdate]) -> None: + builtin_values = { + update.field_name: update.value + for update in updates + if update.custom_field is None + } + original_values = {name: getattr(document, name) for name in builtin_values} + try: + with transaction.atomic(): + for update in updates: + if update.custom_field is not None: + CustomFieldInstance.objects.update_or_create( + document=document, + field=update.custom_field, + defaults={update.field_name: update.value}, + ) + if builtin_values: + for name, value in builtin_values.items(): + setattr(document, name, value) + document.save(update_fields=list(builtin_values)) + except Exception: + # Database rollback does not restore attributes on the model instance. + for name, value in original_values.items(): + setattr(document, name, value) + raise def _convert_value(value: str, data_type: str) -> object | None: