🐛 fix method naming

Author: ianardee

mindee/image/extracted_image.py

@@ -48,7 +48,7 @@ def __init__(
         self._element_id = 0 if element_id is None else element_id
 
     @requires_pillow
-    def save_to_file(self, output_path: Path | str):
+    def write_to_file(self, output_path: Path | str):
         """
         Saves the document to a file.
 

mindee/image/extracted_images.py

@@ -9,4 +9,4 @@ class ExtractedImages(list[ExtractedImage]):
     def save_all_to_disk(self, output_path: Path | str) -> None:
         """Save all extracted images to disk."""
         for image in self:
-            image.save_to_file(output_path)
+            image.write_to_file(output_path)

mindee/pdf/extracted_pdf.py

@@ -3,44 +3,27 @@
 from pathlib import Path
 from typing import BinaryIO
 
-from mindee.dependencies.checkers import PYPDFIUM2_AVAILABLE
-from mindee.dependencies.decorators import requires_pypdfium2
 from mindee.error.mindee_error import MindeeError
 from mindee.input.bytes_input import BytesInput
 
-if PYPDFIUM2_AVAILABLE:
-    # pylint: disable=import-error
-    import pypdfium2 as pdfium
-else:
-    pdfium = None  # pylint: disable=invalid-name
-
 
 class ExtractedPDF:
     """An extracted sub-Pdf."""
 
     buffer: BinaryIO
+    """PDF content as a byte stream."""
     filename: str
-    _page_indexes: tuple[int, int]
+    """Name of the file when writing to disk."""
+    _page_range: tuple[int, int]
 
     def __init__(
-        self, pdf_byte_stream: BinaryIO, filename: str, page_indexes: tuple[int, int]
+        self, pdf_byte_stream: BinaryIO, filename: str, page_range: tuple[int, int]
     ):
         self.buffer = pdf_byte_stream
         self.filename = filename
-        self._page_indexes = page_indexes
-
-    @requires_pypdfium2
-    def get_page_count(self) -> int:
-        """Get the number of pages in the PDF file."""
-        try:
-            pdf = pdfium.PdfDocument(self.buffer)
-            return len(pdf)
-        except Exception as e:
-            raise MindeeError(
-                "Could not retrieve page count from Extracted PDF object."
-            ) from e
+        self._page_range = page_range
 
-    def save_to_file(self, output_path: Path | str):
+    def write_to_file(self, output_path: Path | str):
         """
         Writes the contents of the current PDF object to a file.
 
@@ -66,6 +49,15 @@ def as_input_source(self) -> BytesInput:
         return BytesInput(self.buffer.read(), self.filename)
 
     @property
-    def page_indexes(self) -> tuple[int, int]:
-        """This PDF was extracted from this page range of the original PDF."""
-        return self._page_indexes
+    def page_range(self) -> tuple[int, int]:
+        """
+        This PDF was extracted from this page range of the original PDF.
+        The first number is the index of the first page.
+        The second number is the index of the last page.
+        """
+        return self._page_range
+
+    @property
+    def page_count(self) -> int:
+        """The number of pages in this PDF file."""
+        return self._page_range[1] - self._page_range[0] + 1

mindee/pdf/extracted_pdfs.py

@@ -10,4 +10,4 @@ def save_all_to_disk(self, output_path: Path | str) -> None:
         """Save all extracted images to disk."""
 
         for image in self:
-            image.save_to_file(output_path)
+            image.write_to_file(output_path)

mindee/pdf/pdf_extractor.py

@@ -28,23 +28,19 @@ class PDFExtractor:
 
     _source_pdf: BinaryIO
     _filename: str
+    _page_count: int
 
     @requires_pillow
     def __init__(self, local_input: LocalInputSource):
         self._filename = local_input.filename
+        self._page_count = local_input.page_count
         if local_input.is_pdf():
             self._source_pdf = local_input.file_object
         else:
             pdf_image = Image.open(local_input.file_object)
             self._source_pdf = io.BytesIO()
             pdf_image.save(self._source_pdf, format="PDF")
 
-    @requires_pypdfium2
-    def get_page_count(self) -> int:
-        """Get the number of pages in the PDF file."""
-        pdf = pdfium.PdfDocument(self._source_pdf)
-        return len(pdf)
-
     @requires_pypdfium2
     def cut_pages(self, page_indexes: list) -> BinaryIO:
         """
@@ -78,7 +74,7 @@ def extract_sub_documents(
             if not page_index_elem or len(page_index_elem) == 0:
                 raise MindeeError("Empty indexes aren't allowed for extraction.")
             for page_index in page_index_elem:
-                if page_index > self.get_page_count():
+                if page_index > self._page_count:
                     raise MindeeError(f"Index {page_index} is out of range.")
             first_page = page_index_elem[0]
             last_page = page_index_elem[len(page_index_elem) - 1]

tests/v1/extraction/test_invoice_splitter_auto_extraction.py

@@ -40,7 +40,7 @@ def test_pdf_should_extract_invoices_strict():
     )
     inference = response.document.inference
     pdf_extractor = PDFExtractor(invoice_splitter_input)
-    assert pdf_extractor.get_page_count() == 2
+    assert pdf_extractor.page_count == 2
 
     extracted_pdfs_not_strict = pdf_extractor.extract_invoices(
         inference.prediction.invoice_page_groups

tests/v1/extraction/test_pdf_extractor.py

@@ -39,7 +39,8 @@ def test_image_should_extract_pdf(invoice_default_sample_path):
     jpg_input = PathInput(invoice_default_sample_path)
     assert not jpg_input.is_pdf()
     extractor = PDFExtractor(jpg_input)
-    assert extractor.get_page_count() == 1
+    extracted_pdfs = extractor.extract_documents([[0]])
+    assert len(extracted_pdfs) == 1
 
 
 @pytest.mark.pillow
@@ -48,20 +49,20 @@ def test_pdf_should_extract_invoices_no_strict(
     invoice_splitter_5p_path, loaded_prediction
 ):
     pdf_input = PathInput(invoice_splitter_5p_path)
+    assert pdf_input.page_count == 5
     extractor = PDFExtractor(pdf_input)
-    assert extractor.get_page_count() == 5
     extracted_pdfs_no_strict = extractor.extract_invoices(
         loaded_prediction.invoice_page_groups
     )
 
     assert len(extracted_pdfs_no_strict) == 3
-    assert extracted_pdfs_no_strict[0].get_page_count() == 1
+    assert extracted_pdfs_no_strict[0].page_count == 1
     assert extracted_pdfs_no_strict[0].filename == "invoice_5p_pages-001-001.pdf"
 
-    assert extracted_pdfs_no_strict[1].get_page_count() == 3
+    assert extracted_pdfs_no_strict[1].page_count == 3
     assert extracted_pdfs_no_strict[1].filename == "invoice_5p_pages-002-004.pdf"
 
-    assert extracted_pdfs_no_strict[2].get_page_count() == 1
+    assert extracted_pdfs_no_strict[2].page_count == 1
     assert extracted_pdfs_no_strict[2].filename == "invoice_5p_pages-005-005.pdf"
 
 
@@ -71,15 +72,16 @@ def test_pdf_should_extract_invoices_strict(
     invoice_splitter_5p_path, loaded_prediction
 ):
     pdf_input = PathInput(invoice_splitter_5p_path)
+    assert pdf_input.page_count == 5
+
     extractor = PDFExtractor(pdf_input)
-    assert extractor.get_page_count() == 5
     extracted_pdfs_strict = extractor.extract_invoices(
         loaded_prediction.invoice_page_groups, True
     )
 
     assert len(extracted_pdfs_strict) == 2
-    assert extracted_pdfs_strict[0].get_page_count() == 1
+    assert extracted_pdfs_strict[0].page_count == 1
     assert extracted_pdfs_strict[0].filename == "invoice_5p_pages-001-001.pdf"
 
-    assert extracted_pdfs_strict[1].get_page_count() == 4
+    assert extracted_pdfs_strict[1].page_count == 4
     assert extracted_pdfs_strict[1].filename == "invoice_5p_pages-002-005.pdf"

tests/v2/file_operations/test_split_operation.py

@@ -32,9 +32,9 @@ def test_default_split():
     extracted_splits = response.inference.result.extract_from_input_source(input_sample)
     assert len(extracted_splits) == 2
 
-    assert extracted_splits[0].get_page_count() == 1
+    assert extracted_splits[0].page_count == 1
     assert extracted_splits[0].filename == "default_sample_pages-001-001.pdf"
-    assert extracted_splits[1].get_page_count() == 1
+    assert extracted_splits[1].page_count == 1
     assert extracted_splits[1].filename == "default_sample_pages-002-002.pdf"
 
 
@@ -46,11 +46,11 @@ def test_multi_page_receipt_split(splits_5p, splits_multi_page_json_path):
     extracted_splits = response.inference.result.extract_from_input_source(input_sample)
     assert len(extracted_splits) == 3
 
-    assert extracted_splits[0].get_page_count() == 1
+    assert extracted_splits[0].page_count == 1
     assert extracted_splits[0].filename == "invoice_5p_pages-001-001.pdf"
-    assert extracted_splits[1].get_page_count() == 3
+    assert extracted_splits[1].page_count == 3
     assert extracted_splits[1].filename == "invoice_5p_pages-002-004.pdf"
-    assert extracted_splits[2].get_page_count() == 1
+    assert extracted_splits[2].page_count == 1
     assert extracted_splits[2].filename == "invoice_5p_pages-005-005.pdf"
 
 
@@ -62,4 +62,4 @@ def test_multi_page_receipt_single_split(splits_5p, splits_multi_page_json_path)
     split = response.inference.result.splits[1]
     extracted_split = split.extract_from_input_source(input_sample)
 
-    assert extracted_split.get_page_count() == 3
+    assert extracted_split.page_count == 3