better page count system

Author: ianardee

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
@@ -107,11 +108,14 @@ def page_count(self) -> int:
 
         :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."""
@@ -177,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:
         """

tests/input/test_apply_page_options.py

@@ -42,6 +42,7 @@ def test_pdf_reconstruct_ok():
 @pytest.mark.parametrize("numb_pages", [1, 2, 3])
 def test_process_pdf_cut_n_pages(numb_pages: int):
     input_source = PathInput(FILE_TYPES_DIR / "pdf" / "multipage.pdf")
+    assert input_source.page_count == 12
     input_source.process_pdf(
         behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[0, -2, -1][:numb_pages]
     )
@@ -52,6 +53,7 @@ def test_process_pdf_cut_n_pages(numb_pages: int):
 @pytest.mark.parametrize("numb_pages", [1, 2, 3])
 def test_apply_pages_pdf_cut_n_pages(numb_pages: int):
     input_source = PathInput(FILE_TYPES_DIR / "pdf" / "multipage.pdf")
+    assert input_source.page_count == 12
     input_source.apply_page_options(
         PageOptions(on_min_pages=2, page_indexes=[0, -2, -1][:numb_pages])
     )
@@ -61,7 +63,7 @@ def test_apply_pages_pdf_cut_n_pages(numb_pages: int):
 
 def test_pdf_keep_5_first_pages():
     input_source = PathInput(FILE_TYPES_DIR / "pdf" / "multipage.pdf")
-    assert input_source.is_pdf() is True
+    assert input_source.page_count == 12
     input_source.process_pdf(
         behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[0, 1, 2, 3, 4]
     )
@@ -70,11 +72,11 @@ def test_pdf_keep_5_first_pages():
 
 def test_pdf_keep_invalid_pages():
     input_source = PathInput(FILE_TYPES_DIR / "pdf" / "multipage.pdf")
-    assert input_source.is_pdf() is True
+    assert input_source.page_count == 12
     input_source.process_pdf(
         behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[0, 1, 17]
     )
-    assert input_source.count_doc_pages() == 2
+    assert input_source.page_count == 2
 
 
 def test_pdf_remove_5_last_pages():
@@ -83,7 +85,7 @@ def test_pdf_remove_5_last_pages():
     input_source.process_pdf(
         behavior=REMOVE, on_min_pages=2, page_indexes=[-5, -4, -3, -2, -1]
     )
-    assert input_source.count_doc_pages() == 7
+    assert input_source.page_count == 7
 
 
 def test_pdf_remove_5_first_pages():
@@ -92,14 +94,14 @@ def test_pdf_remove_5_first_pages():
     input_source.process_pdf(
         behavior=REMOVE, on_min_pages=2, page_indexes=list(range(5))
     )
-    assert input_source.count_doc_pages() == 7
+    assert input_source.page_count == 7
 
 
 def test_pdf_remove_invalid_pages():
     input_source = PathInput(FILE_TYPES_DIR / "pdf" / "multipage.pdf")
     assert input_source.is_pdf() is True
     input_source.process_pdf(behavior=REMOVE, on_min_pages=2, page_indexes=[16])
-    assert input_source.count_doc_pages() == 12
+    assert input_source.page_count == 12
 
 
 def test_pdf_keep_no_pages():
@@ -129,23 +131,23 @@ def test_pdf_input_from_file():
         input_source = FileInput(fp)
         assert input_source.is_pdf() is True
         input_source.process_pdf(behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[0])
-    assert input_source.count_doc_pages() == 1
+    assert input_source.page_count == 1
 
 
 def test_pdf_input_from_base64():
     with open(PRODUCT_DATA_DIR / "invoices" / "invoice_10p.txt", "rt") as fp:
         input_source = Base64Input(fp.read(), filename="invoice_10p.pdf")
     assert input_source.is_pdf() is True
     input_source.process_pdf(behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[0])
-    assert input_source.count_doc_pages() == 1
+    assert input_source.page_count == 1
 
 
 def test_pdf_input_from_bytes():
     with open(PRODUCT_DATA_DIR / "invoices" / "invoice_10p.pdf", "rb") as fp:
         input_source = BytesInput(fp.read(), filename="invoice_10p.pdf")
     assert input_source.is_pdf() is True
     input_source.process_pdf(behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[0])
-    assert input_source.count_doc_pages() == 1
+    assert input_source.page_count == 1
 
 
 def test_pdf_blank_check():
@@ -158,4 +160,4 @@ def test_pdf_blank_check():
         input_source.process_pdf(behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[0])
 
     input_not_blank = PathInput(FILE_TYPES_DIR / "pdf" / "not_blank_image_only.pdf")
-    assert input_not_blank.count_doc_pages() == 1
+    assert input_not_blank.page_count == 1

tests/input/test_inputs.py

@@ -8,6 +8,7 @@
     Base64Input,
     BytesInput,
     FileInput,
+    LocalInputSource,
     PathInput,
     UrlInputSource,
 )
@@ -57,42 +58,37 @@ def test_pdf_input_from_url():
 )
 
 
-@pytest.mark.parametrize(("filename", "mimetype"), TEST_IMAGES)
-def test_image_input_from_path(filename, mimetype):
-    input_source = PathInput(FILE_TYPES_DIR / filename)
+def _assert_image(input_source: LocalInputSource, mimetype: str) -> None:
     assert input_source.file_mimetype == mimetype
     assert input_source.is_pdf() is False
     assert input_source.page_count == 1
-    assert isinstance(input_source.file_object, io.BufferedReader)
+    assert isinstance(input_source.file_object.read(15), bytes)
+
+
+@pytest.mark.parametrize(("filename", "mimetype"), TEST_IMAGES)
+def test_image_input_from_path(filename, mimetype):
+    input_source = PathInput(FILE_TYPES_DIR / filename)
+    _assert_image(input_source, mimetype)
 
 
 @pytest.mark.parametrize(("filename", "mimetype"), TEST_IMAGES)
 def test_image_input_from_file(filename, mimetype):
     with open(FILE_TYPES_DIR / filename, "rb") as fp:
         input_source = FileInput(fp)
-        assert input_source.file_mimetype == mimetype
-        assert input_source.is_pdf() is False
-        assert input_source.page_count == 1
-        assert isinstance(input_source.file_object, io.BufferedReader)
+        _assert_image(input_source, mimetype)
 
 
 @pytest.mark.parametrize(("filename", "mimetype"), TEST_IMAGES)
 def test_image_input_from_bytes(filename, mimetype):
     file_bytes = open(FILE_TYPES_DIR / filename, "rb").read()
     input_source = BytesInput(file_bytes, filename=filename)
-    assert input_source.file_mimetype == mimetype
-    assert input_source.is_pdf() is False
-    assert input_source.page_count == 1
-    assert isinstance(input_source.file_object, io.BytesIO)
+    _assert_image(input_source, mimetype)
 
 
 def test_image_input_from_base64():
     base64_input = open(FILE_TYPES_DIR / "receipt.txt", "r").read()
     input_source = Base64Input(base64_input, filename="receipt.jpg")
-    assert input_source.file_mimetype == "image/jpeg"
-    assert input_source.is_pdf() is False
-    assert input_source.page_count == 1
-    assert isinstance(input_source.file_object, io.BytesIO)
+    _assert_image(input_source, mimetype="image/jpeg")
 
 
 def test_txt_input_from_path():