rework again

Author: ianardee

mindee/pdf/pdf_extractor.py

@@ -42,23 +42,36 @@ 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_range=(first_page, last_page),
+        )
 
     @requires_pypdfium2
-    def extract_sub_documents(
+    def extract_multiple_documents(
         self, page_indexes: list[list[int]]
     ) -> list[ExtractedPDF]:
         """
@@ -67,35 +80,13 @@ def extract_sub_documents(
         :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)
+            extracted_pdfs.append(self.extract_single_document(page_index_elem))
         return 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

@@ -35,4 +35,4 @@ def extract_multiple_splits(
         page_groups.append(list(range(split[0], split[1] + 1)))
     if len(splits) < 1:
         raise MindeeError("No indexes provided.")
-    return ExtractedPDFs(pdf_extractor.extract_sub_documents(page_groups))
+    return ExtractedPDFs(pdf_extractor.extract_multiple_documents(page_groups))

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_range == (0, 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