Merge pull request #122 from zane-programs/main

Add show_progress option to to_markdown()

Author: JorjMcKie

pymupdf4llm/pymupdf4llm/helpers/progress.py

@@ -0,0 +1,64 @@
+import sys
+from typing import List, Any
+
+
+class _ProgressBar:
+    def __init__(self, items: List[Any], progress_width: int = 40):
+        self._len = len(items)
+        self._iter = iter(items)
+        self._len_digits = len(str(self._len))
+        self._progress_width = progress_width
+        self._progress_bar = 0
+        self._current_index = 0
+
+        # Calculate the increment for each item based on the list length and the progress width
+        self._increment = self._progress_width / self._len if self._len else 1
+
+        # Init progress bar
+        sys.stdout.write("[%s] (0/%d)" %
+                         (" " * self._progress_width, self._len))
+        sys.stdout.flush()
+        sys.stdout.write(
+            "\b" * (self._progress_width + len(str(self._len)) + 6))
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            result = next(self._iter)
+        except StopIteration as e:
+            # End progress on StopIteration
+            sys.stdout.write("]\n")
+            raise e
+
+        # Update the current index
+        self._current_index += 1
+
+        # Add the increment to the progress bar and calculate how many "=" to add
+        self._progress_bar += self._increment
+        while self._progress_bar >= 1:
+            sys.stdout.write("=")
+            sys.stdout.flush()
+            self._progress_bar -= 1
+
+        # Update the numerical progress
+        padded_index = str(self._current_index).rjust(self._len_digits)
+        progress_info = f" ({padded_index}/{self._len})"
+        sys.stdout.write(
+            "\b" * (self._progress_width + len(progress_info) + 1))
+        sys.stdout.write("[")
+        sys.stdout.write("=" * int(self._current_index *
+                         self._progress_width / self._len))
+        sys.stdout.write(" " * (self._progress_width -
+                         int(self._current_index * self._progress_width / self._len)))
+        sys.stdout.write("]" + progress_info)
+        sys.stdout.flush()
+        sys.stdout.write("\b" * (self._progress_width - int(self._current_index * self._progress_width / self._len)
+                         + len(progress_info) + 1))
+
+        return result
+
+
+def ProgressBar(list: List[Any], progress_width: int = 40):
+    return iter(_ProgressBar(list, progress_width))

pymupdf4llm/pymupdf4llm/helpers/pymupdf_rag.py

@@ -33,6 +33,7 @@
 
 from pymupdf4llm.helpers.get_text_lines import get_raw_lines, is_white
 from pymupdf4llm.helpers.multi_column import column_boxes
+from pymupdf4llm.helpers.progress import ProgressBar
 
 bullet = tuple(
     [
@@ -227,6 +228,7 @@ def to_markdown(
     graphics_limit=None,
     ignore_code=False,
     extract_words=False,
+    show_progress=True,
 ) -> str:
     """Process the document and return the text of the selected pages.
 
@@ -245,6 +247,7 @@ def to_markdown(
         page_height: (float) assumption if page layout is variable.
         table_strategy: choose table detection strategy
         graphics_limit: (int) ignore page with too many vector graphics.
+        show_progress: (bool) print progress as each page is processed.
 
     """
     if write_images is False and force_text is False:
@@ -795,6 +798,9 @@ def get_page_output(doc, pno, margins, textflags):
     # read the Table of Contents
     toc = doc.get_toc()
     textflags = pymupdf.TEXT_MEDIABOX_CLIP
+    if show_progress:
+        print(f"Processing {doc.name}...")
+        pages = ProgressBar(pages)
     for pno in pages:
         page_output, images, tables, graphics, words = get_page_output(
             doc, pno, margins, textflags