pr comments

Author: aadland6

test_unstructured/partition/pdf_image/test_pdf.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import base64
+import io
 import logging
 import math
 import os
@@ -1542,6 +1543,75 @@ def test_pdf_hi_res_max_pages_argument(filename, pdf_hi_res_max_pages, expected_
             )
 
 
+def test_is_pdf_too_complex_skips_small_file_size():
+    assert not pdf.is_pdf_too_complex(file=b"tiny", min_file_size_bytes=10)
+
+
+def test_is_pdf_too_complex_detects_vector_heavy_page():
+    class MockStream:
+        def get_data(self):
+            return b" ".join([b"m"] * 120 + [b"Tj"] * 2)
+
+    reader = mock.Mock()
+    reader.pages = [{"/Contents": MockStream()}]
+
+    with mock.patch.object(pdf, "PdfReader", return_value=reader):
+        assert pdf.is_pdf_too_complex(
+            file=b"x" * 20,
+            max_graphics_ops=100,
+            min_graphics_to_text_ratio=20.0,
+            min_file_size_bytes=1,
+            min_raw_stream_bytes=1,
+        )
+
+
+def test_is_pdf_too_complex_skips_pages_without_contents():
+    reader = mock.Mock()
+    reader.pages = [{"/Contents": None}]
+
+    with mock.patch.object(pdf, "PdfReader", return_value=reader):
+        assert not pdf.is_pdf_too_complex(
+            file=b"x" * 20,
+            min_file_size_bytes=1,
+            min_raw_stream_bytes=1,
+        )
+
+
+def test_is_pdf_too_complex_skips_small_content_streams():
+    class MockStream:
+        def get_data(self):
+            return b"m Tj"
+
+    reader = mock.Mock()
+    reader.pages = [{"/Contents": MockStream()}]
+
+    with mock.patch.object(pdf, "PdfReader", return_value=reader):
+        assert not pdf.is_pdf_too_complex(
+            file=b"x" * 20,
+            max_graphics_ops=1,
+            min_graphics_to_text_ratio=1.0,
+            min_file_size_bytes=1,
+            min_raw_stream_bytes=20,
+        )
+
+
+def test_is_pdf_too_complex_restores_file_cursor_position():
+    file = io.BytesIO(b"x" * 20)
+    file.seek(7)
+
+    reader = mock.Mock()
+    reader.pages = []
+
+    with mock.patch.object(pdf, "PdfReader", return_value=reader):
+        assert not pdf.is_pdf_too_complex(
+            file=file,
+            min_file_size_bytes=1,
+            min_raw_stream_bytes=1,
+        )
+
+    assert file.tell() == 7
+
+
 def test_document_to_element_list_omits_coord_system_when_coord_points_absent():
     # TODO (yao): investigate why we need this test. The LayoutElement definition suggests bbox
     # can't be None and it has to be a Rectangle object that has x1, y1, x2, y2 attributes.

unstructured/partition/pdf.py

@@ -7,7 +7,7 @@
 import re
 import warnings
 from pathlib import Path
-from typing import IO, TYPE_CHECKING, Any, Optional, cast
+from typing import IO, TYPE_CHECKING, Any, Optional, Union, cast
 
 import numpy as np
 import wrapt
@@ -103,6 +103,8 @@
     rb"(?:^|(?<=\s))" rb"(?:Tj|TJ|'|\"|Tf|Td|TD|Tm|T\*|BT|ET)" rb"(?=\s|$)",
     re.MULTILINE,
 )
+DEFAULT_MIN_FILE_SIZE_BYTES = 1 * 1024 * 1024  # 1 MB
+DEFAULT_MIN_RAW_STREAM_BYTES = 100_000  # 100 KB
 
 # increase the max pixels so high dpi values like 300 can still be under the PIL limit
 PILImage.MAX_IMAGE_PIXELS = 5e8
@@ -593,11 +595,11 @@ def check_pdf_hi_res_max_pages_exceeded(
 
 def is_pdf_too_complex(
     filename: str = "",
-    file: Optional[bytes | IO[bytes]] = None,
+    file: Optional[Union[bytes, IO[bytes]]] = None,
     max_graphics_ops: int = 10_000,
     min_graphics_to_text_ratio: float = 20.0,
-    min_file_size_bytes: int = int(1 * 1024 * 1024),  # 1 MB
-    min_raw_stream_bytes: int = 100_000,
+    min_file_size_bytes: int = DEFAULT_MIN_FILE_SIZE_BYTES,
+    min_raw_stream_bytes: int = DEFAULT_MIN_RAW_STREAM_BYTES,
 ) -> bool:
     """Check if a PDF is likely a complex vector drawing (e.g., CAD/engineering docs)
     that would be extremely slow or produce garbage results with PDFMiner text extraction.
@@ -624,7 +626,7 @@ def is_pdf_too_complex(
         Minimum ratio of graphics ops to text ops required (in conjunction with
         `max_graphics_ops`) to flag a page as too complex.
     min_file_size_bytes
-        Skip the complexity check entirely for files smaller than this (default 2 MB).
+        Skip the complexity check entirely for files smaller than this (default 1 MB).
     min_raw_stream_bytes
         Skip operator counting for pages whose decoded content stream is smaller than
         this (default 100 KB). Small streams can't have enough operators to trigger