♻️ harmonize getting page count from a local input source (#364)

Author: ianardee

.pre-commit-config.yaml

@@ -27,7 +27,7 @@ repos:
       - id: gitleaks
 
   - repo: https://github.com/PyCQA/pylint
-    rev: v3.3.1
+    rev: v3.3.9
     hooks:
       - id: pylint
         name: pylint

examples/auto_invoice_splitter_extraction_example.py

@@ -11,7 +11,7 @@
 def parse_invoice(file_path):
     input_source = PathInput(file_path)
 
-    if input_source.is_pdf() and input_source.count_doc_pages() > 1:
+    if input_source.is_pdf() and input_source.page_count > 1:
         parse_multi_page(input_source)
     else:
         parse_single_page(input_source)

mindee/extraction/multi_receipts_extractor/multi_receipts_extractor.py

@@ -24,7 +24,7 @@ def extract_receipts(
         raise MindeeError(
             "No possible receipts candidates found for MultiReceipts extraction."
         )
-    for page_id in range(input_source.count_doc_pages()):
+    for page_id in range(input_source.page_count):
         receipt_positions = [
             receipt.bounding_box
             for receipt in inference.pages[page_id].prediction.receipts

mindee/input/sources/local_input_source.py

@@ -36,6 +36,7 @@ class LocalInputSource:
     file_mimetype: str
     input_type: InputType
     filepath: Optional[str]
+    _page_count: Optional[int] = None
 
     def __init__(self, input_type: InputType):
         self.input_type = input_type
@@ -100,17 +101,25 @@ def is_pdf(self) -> bool:
         """:return: True if the file is a PDF."""
         return self.file_mimetype == "application/pdf"
 
-    def count_doc_pages(self) -> int:
+    @property
+    def page_count(self) -> int:
         """
-        Count the pages in the PDF.
+        Count the pages in the document.
 
-        :return: the number of pages.
+        :return: The number of pages.
         """
-        if self.is_pdf():
-            self.file_object.seek(0)
-            pdf = pdfium.PdfDocument(self.file_object)
-            return len(pdf)
-        return 1
+        if self._page_count is None:
+            if self.is_pdf():
+                self.file_object.seek(0)
+                pdf = pdfium.PdfDocument(self.file_object)
+                self._page_count = len(pdf)
+            else:
+                self._page_count = 1
+        return self._page_count
+
+    def count_doc_pages(self) -> int:
+        """Deprecated. Use ``page_count`` instead."""
+        return self.page_count
 
     def apply_page_options(self, page_options: PageOptions) -> None:
         """Apply cut and merge options on multipage documents."""
@@ -131,10 +140,10 @@ def process_pdf(
         """Run any required processing on a PDF file."""
         if self.is_pdf_empty():
             raise MindeeSourceError(f"PDF pages are empty in: {self.filename}")
-        pages_count = self.count_doc_pages()
-        if on_min_pages > pages_count:
+        page_count = self.page_count
+        if on_min_pages > page_count:
             return
-        all_pages = list(range(pages_count))
+        all_pages = list(range(page_count))
         if behavior == KEEP_ONLY:
             pages_to_keep = set()
             for page_id in page_indexes:
@@ -161,7 +170,7 @@ def merge_pdf_pages(self, page_numbers: set) -> None:
         """
         Create a new PDF from pages and set it to ``file_object``.
 
-        :param page_numbers: List of pages number to use for merging in the original PDF.
+        :param page_numbers: List of page numbers to use for merging in the original PDF.
         :return: None
         """
         self.file_object.seek(0)
@@ -172,6 +181,7 @@ def merge_pdf_pages(self, page_numbers: set) -> None:
         bytes_io = io.BytesIO()
         new_pdf.save(bytes_io)
         self.file_object = bytes_io
+        self._page_count = len(new_pdf)
 
     def is_pdf_empty(self) -> bool:
         """

pyproject.toml

@@ -44,15 +44,15 @@ Changelog = "https://github.com/mindee/mindee-api-python/blob/main/CHANGELOG.md"
 
 [project.optional-dependencies]
 lint = [
-  "pylint==3.3.1",
-  "pre-commit~=3.2.2",
-  "types-pytz>=2023.3",
+  "pylint==3.3.9",
+  "pre-commit~=3.6.0",
+  "types-pytz>=2024.2",
   "types-requests>=2.31",
 ]
 test = [
   "toml~=0.10.2",
   "pytest~=7.4",
-  "pytest-cov~=4.1",
+  "pytest-cov~=5.0",
 ]
 docs = [
   "sphinx~=5.3",

tests/extraction/test_image_extractor.py

@@ -6,7 +6,7 @@
 from mindee.extraction.common.image_extractor import extract_multiple_images_from_source
 from mindee.input.sources.path_input import PathInput
 from mindee.product.barcode_reader.barcode_reader_v1 import BarcodeReaderV1
-from tests.test_inputs import PRODUCT_DATA_DIR
+from tests.utils import PRODUCT_DATA_DIR
 
 
 @pytest.fixture

tests/extraction/test_invoice_splitter_auto_extraction.py

@@ -9,8 +9,7 @@
 from mindee.product.invoice.invoice_v4 import InvoiceV4
 from mindee.product.invoice_splitter.invoice_splitter_v1 import InvoiceSplitterV1
 from tests.product import get_id, get_version
-from tests.test_inputs import PRODUCT_DATA_DIR
-from tests.utils import levenshtein_ratio
+from tests.utils import PRODUCT_DATA_DIR, levenshtein_ratio
 
 
 @pytest.fixture

tests/extraction/test_multi_receipts_extractor.py

@@ -10,7 +10,7 @@
 from mindee.product.multi_receipts_detector.multi_receipts_detector_v1 import (
     MultiReceiptsDetectorV1,
 )
-from tests.test_inputs import PRODUCT_DATA_DIR
+from tests.utils import PRODUCT_DATA_DIR
 
 
 @pytest.fixture

tests/extraction/test_pdf_extractor.py

@@ -8,7 +8,7 @@
 from mindee.product.invoice_splitter.invoice_splitter_v1_document import (
     InvoiceSplitterV1Document,
 )
-from tests.test_inputs import PRODUCT_DATA_DIR
+from tests.utils import PRODUCT_DATA_DIR
 
 
 @pytest.fixture

tests/extras/test_extras_integration.py

@@ -3,7 +3,7 @@
 from mindee import Client
 from mindee.product.international_id.international_id_v2 import InternationalIdV2
 from mindee.product.invoice.invoice_v4 import InvoiceV4
-from tests.product import PRODUCT_DATA_DIR
+from tests.utils import PRODUCT_DATA_DIR
 
 
 @pytest.fixture