feat: report page progress from partition_pdf and partition_image

RU
2d ago
+588-46

Summary

Partitioning a long PDF with partition_pdf() or partition_image() can take minutes with no sign of how far along it is (#2351). This adds two opt-in arguments, both off by default, so callers can see progress without changing anything else:

  • progress_callback: called with a PartitionProgress(stage, pages_done, total_pages) each time a page finishes. The caller decides what to do with it (log, update a UI, feed a job queue).
  • show_progress: draws a tqdm bar on stderr for each stage. tqdm is already a core dependency, so no new requirement.

Where the pages are counted:

  • unstructured/partition/utils/progress.py (new): PartitionProgress, the report_page_progress() context manager and the reporter. When neither argument is set it yields None, so the hot loops only pay for one if reporter: check.
  • unstructured/partition/pdf.py: both arguments are explicit parameters on partition_pdf() and partition_pdf_or_image() (so they do not leak through **kwargs). The pdfminer text pass in _process_pdfminer_pages() reports stage "text"; the ocr_only loop in _partition_pdf_or_image_with_ocr() reports stage "ocr"; hi_res forwards the arguments to the OCR helpers. The page count for total_pages is read with pypdf only when progress was requested, and a document whose count cannot be read still reports, with total_pages=None.
  • unstructured/partition/pdf_image/ocr.py: process_file_with_ocr() and process_data_with_ocr() accept the same arguments and report stage "ocr" per page for both PDFs and (multi-frame) images.

auto and hi_res can run both passes on one document, so a callback may see the page count restart with a new stage; this is documented in the partition_pdf() docstring and the module docstring. The layout-model phase of hi_res is not reported per page.

Example:

from unstructured.partition.pdf import partition_pdf

def on_progress(p):
    print(f"{p.stage}: {p.pages_done}/{p.total_pages}")

elements = partition_pdf("long.pdf", progress_callback=on_progress)
# or simply
elements = partition_pdf("long.pdf", strategy="hi_res", show_progress=True)

Closes #2351

Validation

  • New test_unstructured/partition/utils/test_progress.py (9 tests): nothing is emitted by default, one event per page with the right stage and counts, unknown totals, the tqdm bar is only drawn when show_progress=True and is closed even when the loop raises.
  • test_unstructured/partition/pdf_image/test_pdf.py: fast on layout-parser-paper-fast.pdf emits one "text" event per page with total_pages equal to the page count (filename and file object); the page count is not read unless progress is requested; an unreadable page count degrades to total_pages=None; hi_res forwards progress_callback / show_progress to process_file_with_ocr and passes None / False by default; ocr_only emits one "ocr" event per rendered page; partition_image emits a single "ocr" event; show_progress=True draws one bar sized to the page count and closes it.
  • test_unstructured/partition/pdf_image/test_ocr.py: process_file_with_ocr emits one "ocr" event per PDF page and per image, nothing by default; process_data_with_ocr forwards both arguments.
  • ruff check and ruff format --check pass on all changed files.
  • The test_progress.py suite was run locally and passes (9 passed). The test_pdf.py and test_ocr.py additions could not be run locally: this machine only has Python 3.9 and the package requires 3.11+ (ParamSpec import in unstructured/telemetry.py), so they rely on CI. I have kept them to the existing mocking patterns in those files.

Release note

  • partition_pdf() and partition_image() accept progress_callback and show_progress to report per-page progress for the text and OCR passes. Both are off by default.

Checklist

  • The PR title is a user-facing release note in conventional-commit form.
  • Tests were added for the new behaviour.
  • Docstrings and type hints were added for every new function and argument.
  • No change to CHANGELOG.md (frozen) or __version__.