rework again

Author: ianardee

mindee/pdf/extracted_pdf.py

@@ -14,14 +14,14 @@ class ExtractedPDF:
     """PDF content as a byte stream."""
     filename: str
     """Name of the file when writing to disk."""
-    _page_range: tuple[int, int]
+    _page_indexes: list[int]
 
     def __init__(
-        self, pdf_byte_stream: BinaryIO, filename: str, page_range: tuple[int, int]
+        self, pdf_byte_stream: BinaryIO, filename: str, page_indexes: list[int]
     ):
         self.buffer = pdf_byte_stream
         self.filename = filename
-        self._page_range = page_range
+        self._page_indexes = page_indexes
 
     def write_to_file(self, output_path: Path | str):
         """
@@ -49,15 +49,13 @@ def as_input_source(self) -> BytesInput:
         return BytesInput(self.buffer.read(), self.filename)
 
     @property
-    def page_range(self) -> tuple[int, int]:
+    def page_indexes(self) -> list[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.
+        0-based indexes of all pages taken from the original PDF.
         """
-        return self._page_range
+        return self._page_indexes
 
     @property
     def page_count(self) -> int:
         """The number of pages in this PDF file."""
-        return self._page_range[1] - self._page_range[0] + 1
+        return len(self._page_indexes)

mindee/pdf/pdf_extractor.py

@@ -9,6 +9,7 @@
 from mindee.error.mindee_error import MindeeError
 from mindee.input.local_input_source import LocalInputSource
 from mindee.pdf.extracted_pdf import ExtractedPDF
+from mindee.pdf.extracted_pdfs import ExtractedPDFs
 
 if PYPDFIUM2_AVAILABLE:
     # pylint: disable=import-error
@@ -42,60 +43,51 @@ def __init__(self, local_input: LocalInputSource):
             pdf_image.save(self._source_pdf, format="PDF")
 
     @requires_pypdfium2
-    def cut_pages(self, page_indexes: list) -> BinaryIO:
+    def extract_single_document(self, page_indexes: list[int]) -> ExtractedPDF:
         """
         Create a new PDF from pages and save it into a buffer.
 
         :param page_indexes: List of pages number to use for merging in the original PDF.
         :return: The buffer containing the new PDF.
         """
+        if not page_indexes or len(page_indexes) == 0:
+            raise MindeeError("Empty indexes aren't allowed for extraction.")
+        for page_index in page_indexes:
+            if page_index > self._page_count:
+                raise MindeeError(f"Index {page_index} is out of range.")
+
         self._source_pdf.seek(0)
         new_pdf = pdfium.PdfDocument.new()
         pdf = pdfium.PdfDocument(self._source_pdf)
         new_pdf.import_pages(pdf, page_indexes)
         bytes_io = io.BytesIO()
         new_pdf.save(bytes_io)
-        return bytes_io
+
+        first_page = page_indexes[0]
+        last_page = page_indexes[len(page_indexes) - 1]
+        return ExtractedPDF(
+            pdf_byte_stream=bytes_io,
+            filename=self._make_filename(first_page, last_page),
+            page_indexes=page_indexes,
+        )
 
     @requires_pypdfium2
-    def extract_sub_documents(
+    def extract_multiple_documents(
         self, page_indexes: list[list[int]]
-    ) -> list[ExtractedPDF]:
+    ) -> ExtractedPDFs:
         """
         Extract the sub-documents from the main pdf, based on the given list of page indexes.
 
         :param page_indexes: 2D list of numbers, representing page indexes.
         :return: A list of created PDFS.
         """
+        if len(page_indexes) < 1:
+            raise MindeeError("No indexes provided.")
         extracted_pdfs: list[ExtractedPDF] = []
-        extension = Path(self._filename).suffix
-        stem = Path(self._filename).stem
         for page_index_elem in page_indexes:
-            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._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]
-            extracted_pdf = ExtractedPDF(
-                self.cut_pages(page_index_elem),
-                f"{stem}_pages-{(first_page + 1):03d}-{(last_page + 1):03d}{extension}",
-                (first_page, last_page),
-            )
-            extracted_pdfs.append(extracted_pdf)
-        return extracted_pdfs
+            extracted_pdfs.append(self.extract_single_document(page_index_elem))
+        return ExtractedPDFs(extracted_pdfs)
 
-    def extract_documents(
-        self,
-        page_indexes: list[list[int]],
-    ) -> list[ExtractedPDF]:
-        """
-        Extracts complete PDFs from the document.
-
-        :param page_indexes: List of sub-lists of pages to keep.
-        :return: A list of extracted invoices.
-        """
-        if len(page_indexes) < 1:
-            raise MindeeError("No indexes provided.")
-        return self.extract_sub_documents(page_indexes)
+    def _make_filename(self, first_page: int, last_page: int) -> str:
+        stem = Path(self._filename).stem
+        return f"{stem}_pages-{(first_page + 1):03d}-{(last_page + 1):03d}.pdf"

mindee/v1/pdf/pdf_extractor.py

@@ -24,11 +24,11 @@ def extract_invoices(
         if len(page_indexes) < 1:
             raise MindeeError("No indexes provided.")
         if not isinstance(page_indexes[0], InvoiceSplitterV1InvoicePageGroup):
-            return self.extract_sub_documents(page_indexes)  # type: ignore
+            return self.extract_multiple_documents(page_indexes)  # type: ignore
 
         if not strict:
             indexes_as_list = [page_index.page_indexes for page_index in page_indexes]  # type: ignore
-            return self.extract_sub_documents(indexes_as_list)
+            return self.extract_multiple_documents(indexes_as_list)
         correct_page_indexes: list[list[int]] = []
         current_list: list[int] = []
         previous_confidence: float | None = None
@@ -49,4 +49,4 @@ def extract_invoices(
                 correct_page_indexes.append(current_list)
                 correct_page_indexes.append(page_list)
             previous_confidence = confidence
-        return self.extract_sub_documents(correct_page_indexes)
+        return self.extract_multiple_documents(correct_page_indexes)

mindee/v2/file_operations/split.py

@@ -15,7 +15,8 @@ def extract_single_split(
     :param split: List of pages to keep.
     :return: Extracted PDF
     """
-    return extract_multiple_splits(input_source, [split])[0]
+    pdf_extractor = PDFExtractor(input_source)
+    return pdf_extractor.extract_single_document(_range_to_indexes(split))
 
 
 def extract_multiple_splits(
@@ -32,7 +33,11 @@ def extract_multiple_splits(
     pdf_extractor = PDFExtractor(input_source)
     page_groups = []
     for split in splits:
-        page_groups.append(list(range(split[0], split[1] + 1)))
+        page_groups.append(_range_to_indexes(split))
     if len(splits) < 1:
         raise MindeeError("No indexes provided.")
-    return ExtractedPDFs(pdf_extractor.extract_sub_documents(page_groups))
+    return pdf_extractor.extract_multiple_documents(page_groups)
+
+
+def _range_to_indexes(split: list[int]) -> list[int]:
+    return list(range(split[0], split[1] + 1))

tests/v1/extraction/test_invoice_splitter_auto_extraction.py

@@ -48,7 +48,7 @@ def test_pdf_should_extract_invoices_strict():
     extracted_pdfs_strict = pdf_extractor.extract_invoices(
         inference.prediction.invoice_page_groups
     )
-    extracted_base_pdfs = pdf_extractor.extract_documents(
+    extracted_base_pdfs = pdf_extractor.extract_multiple_documents(
         [int_list.page_indexes for int_list in inference.prediction.invoice_page_groups]
     )
     for i, extracted_pdf in enumerate(extracted_base_pdfs):

tests/v1/extraction/test_pdf_extractor.py

@@ -8,7 +8,7 @@
 from mindee.v1.product.invoice_splitter.invoice_splitter_v1_document import (
     InvoiceSplitterV1Document,
 )
-from tests.utils import V1_PRODUCT_DATA_DIR
+from tests.utils import OUTPUT_DIR, V1_PRODUCT_DATA_DIR
 
 
 @pytest.fixture
@@ -39,8 +39,12 @@ 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)
-    extracted_pdfs = extractor.extract_documents([[0]])
-    assert len(extracted_pdfs) == 1
+    extracted_pdf = extractor.extract_single_document([0])
+    assert extracted_pdf.page_count == 1
+    assert extracted_pdf.page_indexes == [0]
+    assert extracted_pdf.filename == "default_sample_pages-001-001.pdf"
+    extracted_pdf.write_to_file(OUTPUT_DIR)
+    assert (OUTPUT_DIR / extracted_pdf.filename).exists()
 
 
 @pytest.mark.pillow

tests/v2/file_operations/test_split_operation.py

@@ -47,10 +47,13 @@ def test_multi_page_receipt_split(splits_5p, splits_multi_page_json_path):
     assert len(extracted_splits) == 3
 
     assert extracted_splits[0].page_count == 1
+    assert extracted_splits[0].page_indexes == [0]
     assert extracted_splits[0].filename == "invoice_5p_pages-001-001.pdf"
     assert extracted_splits[1].page_count == 3
+    assert extracted_splits[1].page_indexes == [1, 2, 3]
     assert extracted_splits[1].filename == "invoice_5p_pages-002-004.pdf"
     assert extracted_splits[2].page_count == 1
+    assert extracted_splits[2].page_indexes == [4]
     assert extracted_splits[2].filename == "invoice_5p_pages-005-005.pdf"
 
 

tests/v2/file_operations/test_split_operation_integration.py

@@ -58,7 +58,7 @@ def test_pdf_should_extract_splits():
     for i in range(len(extracted_splits)):
         local_input = PathInput(OUTPUT_DIR / output_files[i])
         try:
-            assert local_input.page_count == extracted_splits[i].get_page_count()
+            assert local_input.page_count == extracted_splits[i].page_count
         finally:
             local_input.close()
     split_input.close()