diff --git a/tests/test_analyze_page.py b/tests/test_analyze_page.py index afcf4ac..f7753ee 100644 --- a/tests/test_analyze_page.py +++ b/tests/test_analyze_page.py @@ -315,3 +315,36 @@ async def test_analyze_page_font_caching(): # But both redactions should share the same font assert result.redactions[0].font is font_match assert result.redactions[1].font is font_match + + +# --------------------------------------------------------------------------- +# test_analyze_page_uses_precomputed_ocr_lines +# --------------------------------------------------------------------------- + +@pytest.mark.anyio +async def test_analyze_page_uses_precomputed_ocr_lines(): + """When ocr_lines is provided, analyze_page should use them directly + without calling ocr_page. The returned lines must be the same object.""" + line = _make_line_from_words(["Hello", "world"]) + precomputed_lines = [line] + page_image = _dummy_page_image() + + mock_ocr = patch( + "unredact.pipeline.analyze_page.ocr_page", + return_value=[], # should never be called + ) + + with ( + mock_ocr as m_ocr, + patch( + "unredact.pipeline.analyze_page.detect_redactions_llm", + new_callable=AsyncMock, + return_value=[], + ), + ): + result = await analyze_page(page_image, ocr_lines=precomputed_lines) + + # ocr_page must not have been called + m_ocr.assert_not_called() + # The returned lines must be the exact same object (not a copy) + assert result.lines is precomputed_lines diff --git a/unredact/pipeline/analyze_page.py b/unredact/pipeline/analyze_page.py index ad65a34..fd617e0 100644 --- a/unredact/pipeline/analyze_page.py +++ b/unredact/pipeline/analyze_page.py @@ -45,6 +45,7 @@ class PageAnalysis: async def analyze_page( page_image: Image.Image, on_progress: callable | None = None, + ocr_lines: list[OcrLine] | None = None, ) -> PageAnalysis: """Run the full per-page analysis pipeline. @@ -61,12 +62,17 @@ async def analyze_page( page_image: PIL Image of the document page. on_progress: Optional callback ``(event_name, data_dict)`` for progress reporting. + ocr_lines: Optional pre-computed OCR lines. When provided the + OCR step is skipped and these lines are used directly. Returns: PageAnalysis with OCR lines and redaction analysis results. """ - # Step 1: OCR (blocking — run in thread) - lines: list[OcrLine] = await asyncio.to_thread(ocr_page, page_image) + # Step 1: OCR (blocking — run in thread), or use pre-computed lines + if ocr_lines is not None: + lines = ocr_lines + else: + lines: list[OcrLine] = await asyncio.to_thread(ocr_page, page_image) if on_progress: on_progress("ocr_done", {"line_count": len(lines)})