Add password with PDF files

Author: pprados

unstructured/partition/image.py

@@ -32,6 +32,7 @@ def partition_image(
     starting_page_number: int = 1,
     extract_forms: bool = False,
     form_extraction_skip_tables: bool = True,
+    password:Optional[str]=None,
     **kwargs: Any,
 ) -> list[Element]:
     """Parses an image into a list of interpreted elements.
@@ -91,6 +92,8 @@ def partition_image(
         (results in adding FormKeysValues elements to output).
     form_extraction_skip_tables
         Whether the form extraction logic should ignore regions designated as Tables.
+    password
+        The password to decrypt the PDF file.
     """
     exactly_one(filename=filename, file=file)
 
@@ -113,5 +116,6 @@ def partition_image(
         starting_page_number=starting_page_number,
         extract_forms=extract_forms,
         form_extraction_skip_tables=form_extraction_skip_tables,
+        password=password,
         **kwargs,
     )

unstructured/partition/pdf.py

@@ -133,6 +133,7 @@ def partition_pdf(
     starting_page_number: int = 1,
     extract_forms: bool = False,
     form_extraction_skip_tables: bool = True,
+    password: Optional[str] = None,
     **kwargs: Any,
 ) -> list[Element]:
     """Parses a pdf document into a list of interpreted elements.
@@ -213,6 +214,7 @@ def partition_pdf(
         starting_page_number=starting_page_number,
         extract_forms=extract_forms,
         form_extraction_skip_tables=form_extraction_skip_tables,
+        password=password,
         **kwargs,
     )
 
@@ -234,6 +236,7 @@ def partition_pdf_or_image(
     starting_page_number: int = 1,
     extract_forms: bool = False,
     form_extraction_skip_tables: bool = True,
+    password: Optional[str] = None,
     **kwargs: Any,
 ) -> list[Element]:
     """Parses a pdf or image document into a list of interpreted elements."""
@@ -262,6 +265,7 @@ def partition_pdf_or_image(
                 languages=languages,
                 metadata_last_modified=metadata_last_modified or last_modified,
                 starting_page_number=starting_page_number,
+                password=password,
                 **kwargs,
             )
             pdf_text_extractable = any(
@@ -311,6 +315,7 @@ def partition_pdf_or_image(
                 starting_page_number=starting_page_number,
                 extract_forms=extract_forms,
                 form_extraction_skip_tables=form_extraction_skip_tables,
+                password=password,
                 **kwargs,
             )
             out_elements = _process_uncategorized_text_elements(elements)
@@ -336,6 +341,7 @@ def partition_pdf_or_image(
                 is_image=is_image,
                 metadata_last_modified=metadata_last_modified or last_modified,
                 starting_page_number=starting_page_number,
+                password=password,
                 **kwargs,
             )
             out_elements = _process_uncategorized_text_elements(elements)
@@ -349,6 +355,7 @@ def extractable_elements(
     languages: Optional[list[str]] = None,
     metadata_last_modified: Optional[str] = None,
     starting_page_number: int = 1,
+    password:Optional[str] = None,
     **kwargs: Any,
 ) -> list[list[Element]]:
     if isinstance(file, bytes):
@@ -359,6 +366,7 @@ def extractable_elements(
         languages=languages,
         metadata_last_modified=metadata_last_modified,
         starting_page_number=starting_page_number,
+        password=password,
         **kwargs,
     )
 
@@ -369,6 +377,7 @@ def _partition_pdf_with_pdfminer(
     languages: list[str],
     metadata_last_modified: Optional[str],
     starting_page_number: int = 1,
+    password:Optional[str] = None,
     **kwargs: Any,
 ) -> list[list[Element]]:
     """Partitions a PDF using PDFMiner instead of using a layoutmodel. Used for faster
@@ -402,6 +411,7 @@ def _partition_pdf_with_pdfminer(
             languages=languages,
             metadata_last_modified=metadata_last_modified,
             starting_page_number=starting_page_number,
+            password=password,
             **kwargs,
         )
 
@@ -416,14 +426,16 @@ def _process_pdfminer_pages(
     metadata_last_modified: Optional[str],
     annotation_threshold: Optional[float] = env_config.PDF_ANNOTATION_THRESHOLD,
     starting_page_number: int = 1,
+    password: Optional[str] = None,
     **kwargs,
 ) -> list[list[Element]]:
     """Uses PDFMiner to split a document into pages and process them."""
 
     elements = []
 
     for page_number, (page, page_layout) in enumerate(
-        open_pdfminer_pages_generator(fp), start=starting_page_number
+        open_pdfminer_pages_generator(fp, password=password),
+            start=starting_page_number,
     ):
         width, height = page_layout.width, page_layout.height
 
@@ -545,6 +557,7 @@ def _partition_pdf_or_image_local(
     extract_forms: bool = False,
     form_extraction_skip_tables: bool = True,
     pdf_hi_res_max_pages: Optional[int] = None,
+    password:Optional[str] = None,
     **kwargs: Any,
 ) -> list[Element]:
     """Partition using package installed locally"""
@@ -586,6 +599,7 @@ def _partition_pdf_or_image_local(
             is_image=is_image,
             model_name=hi_res_model_name,
             pdf_image_dpi=pdf_image_dpi,
+            password=password,
         )
 
         if hi_res_model_name.startswith("chipper"):
@@ -594,7 +608,8 @@ def _partition_pdf_or_image_local(
             final_document_layout = inferred_document_layout
         else:
             extracted_layout = (
-                process_file_with_pdfminer(filename=filename, dpi=pdf_image_dpi)
+                process_file_with_pdfminer(filename=filename, dpi=pdf_image_dpi,
+                                           password=password)
                 if pdf_text_extractable
                 else []
             )
@@ -634,13 +649,15 @@ def _partition_pdf_or_image_local(
                 ocr_mode=ocr_mode,
                 pdf_image_dpi=pdf_image_dpi,
                 ocr_layout_dumper=ocr_layout_dumper,
+                password=password,
             )
     else:
         inferred_document_layout = process_data_with_model(
             file,
             is_image=is_image,
             model_name=hi_res_model_name,
             pdf_image_dpi=pdf_image_dpi,
+            password=password,
         )
 
         if hi_res_model_name.startswith("chipper"):
@@ -652,7 +669,8 @@ def _partition_pdf_or_image_local(
                 file.seek(0)
 
             extracted_layout = (
-                process_data_with_pdfminer(file=file, dpi=pdf_image_dpi)
+                process_data_with_pdfminer(file=file, dpi=pdf_image_dpi,
+                                           password=password)
                 if pdf_text_extractable
                 else []
             )
@@ -694,6 +712,7 @@ def _partition_pdf_or_image_local(
                 ocr_mode=ocr_mode,
                 pdf_image_dpi=pdf_image_dpi,
                 ocr_layout_dumper=ocr_layout_dumper,
+                password=password,
             )
 
     # NOTE(alan): starting with v2, chipper sorts the elements itself.
@@ -734,6 +753,7 @@ def _partition_pdf_or_image_local(
             pdf_image_dpi=pdf_image_dpi,
             extract_image_block_to_payload=extract_image_block_to_payload,
             output_dir_path=extract_image_block_output_dir,
+            password=password,
         )
 
     for el_type in extract_image_block_types:
@@ -807,6 +827,7 @@ def _partition_pdf_or_image_local(
             draw_caption=env_config.ANALYSIS_BBOX_DRAW_CAPTION,
             resize=env_config.ANALYSIS_BBOX_RESIZE,
             format=env_config.ANALYSIS_BBOX_FORMAT,
+            password=password,
         )
 
     return out_elements
@@ -845,6 +866,7 @@ def _partition_pdf_or_image_with_ocr(
     is_image: bool = False,
     metadata_last_modified: Optional[str] = None,
     starting_page_number: int = 1,
+    password:Optional[str] = None,
     **kwargs: Any,
 ):
     """Partitions an image or PDF using OCR. For PDFs, each page is converted
@@ -869,7 +891,8 @@ def _partition_pdf_or_image_with_ocr(
             elements.extend(page_elements)
     else:
         for page_number, image in enumerate(
-            convert_pdf_to_images(filename, file), start=starting_page_number
+            convert_pdf_to_images(filename, file, password=password),
+                start=starting_page_number
         ):
             page_elements = _partition_pdf_or_image_with_ocr_from_image(
                 image=image,

unstructured/partition/pdf_image/analysis/bbox_visualisation.py

@@ -546,6 +546,7 @@ def __init__(
         draw_grid: bool = False,
         resize: Optional[float] = None,
         format: str = "png",
+        password: Optional[str] = None,
     ):
         self.draw_caption = draw_caption
         self.draw_grid = draw_grid
@@ -554,6 +555,7 @@ def __init__(
         self.format = format
         self.drawers = []
         self.file = file
+        self.password = password
 
         super().__init__(filename, save_dir)
 
@@ -678,6 +680,7 @@ def load_source_image(self) -> Generator[Image.Image, None, None]:
                         file=self.file,
                         output_folder=temp_dir,
                         path_only=True,
+                        password=self.password,
                     )
                 except Exception as ex:  # noqa: E722
                     print(

unstructured/partition/pdf_image/analysis/tools.py

@@ -66,6 +66,7 @@ def save_analysis_artifiacts(
     draw_caption: bool = True,
     resize: Optional[float] = None,
     format: str = "png",
+    password: Optional[str] = None,
 ):
     """Save the analysis artifacts for a given file. Loads some settings from
     the environment configuration.
@@ -109,6 +110,7 @@ def save_analysis_artifiacts(
             draw_caption=draw_caption,
             resize=resize,
             format=format,
+            password=password,
         )
 
         for layout_dumper in layout_dumpers:
@@ -125,6 +127,7 @@ def render_bboxes_for_file(
     draw_caption: bool = True,
     resize: Optional[float] = None,
     format: str = "png",
+    password: Optional[str] = None,
 ):
     """Render the bounding boxes for a given layout dimp file.
     To be used for analysis after the partition is performed for
@@ -183,6 +186,7 @@ def render_bboxes_for_file(
             draw_caption=draw_caption,
             resize=resize,
             format=format,
+            passwork=password,
         )
 
         for drawer in layout_drawers:

unstructured/partition/pdf_image/ocr.py

@@ -37,6 +37,7 @@ def process_data_with_ocr(
     ocr_mode: str = OCRMode.FULL_PAGE.value,
     pdf_image_dpi: int = 200,
     ocr_layout_dumper: Optional[OCRLayoutDumper] = None,
+    password: Optional[str] = None,
 ) -> "DocumentLayout":
     """
     Process OCR data from a given data and supplement the output DocumentLayout
@@ -64,6 +65,8 @@ def process_data_with_ocr(
 
     - ocr_layout_dumper (OCRLayoutDumper, optional): The OCR layout dumper to save the OCR layout.
 
+    - password (optional): The password to decrypt the PDF file.
+
     Returns:
         DocumentLayout: The merged layout information obtained after OCR processing.
     """
@@ -84,6 +87,7 @@ def process_data_with_ocr(
             ocr_mode=ocr_mode,
             pdf_image_dpi=pdf_image_dpi,
             ocr_layout_dumper=ocr_layout_dumper,
+            password=password,
         )
 
     return merged_layouts
@@ -100,6 +104,7 @@ def process_file_with_ocr(
     ocr_mode: str = OCRMode.FULL_PAGE.value,
     pdf_image_dpi: int = 200,
     ocr_layout_dumper: Optional[OCRLayoutDumper] = None,
+    password: Optional[str] = None,
 ) -> "DocumentLayout":
     """
     Process OCR data from a given file and supplement the output DocumentLayout
@@ -124,6 +129,8 @@ def process_file_with_ocr(
 
     - pdf_image_dpi (int, optional): DPI (dots per inch) for processing PDF images. Defaults to 200.
 
+    - password (optional): The password to decrypt the PDF file.
+
     Returns:
         DocumentLayout: The merged layout information obtained after OCR processing.
     """
@@ -157,6 +164,7 @@ def process_file_with_ocr(
                     dpi=pdf_image_dpi,
                     output_folder=temp_dir,
                     paths_only=True,
+                    userpw=password,
                 )
                 image_paths = cast(List[str], _image_paths)
                 for i, image_path in enumerate(image_paths):

unstructured/partition/pdf_image/pdf_image_utils.py

@@ -58,6 +58,7 @@ def convert_pdf_to_image(
     dpi: int = 200,
     output_folder: Optional[Union[str, PurePath]] = None,
     path_only: bool = False,
+    password:Optional[str] = None,
 ) -> Union[List[Image.Image], List[str]]:
     """Get the image renderings of the pdf pages using pdf2image"""
 
@@ -71,6 +72,7 @@ def convert_pdf_to_image(
             dpi=dpi,
             output_folder=output_folder,
             paths_only=path_only,
+            userpw=password,
         )
     else:
         images = pdf2image.convert_from_path(
@@ -125,6 +127,7 @@ def save_elements(
     is_image: bool = False,
     extract_image_block_to_payload: bool = False,
     output_dir_path: str | None = None,
+    password:Optional[str] = None,
 ):
     """
     Saves specific elements from a PDF as images either to a directory or embeds them in the
@@ -167,6 +170,7 @@ def save_elements(
                 pdf_image_dpi,
                 output_folder=temp_dir,
                 path_only=True,
+                password=password,
             )
             image_paths = cast(List[str], _image_paths)
 
@@ -389,15 +393,16 @@ def convert_pdf_to_images(
     filename: str = "",
     file: Optional[bytes | IO[bytes]] = None,
     chunk_size: int = 10,
+    password:Optional[str] = None,
 ) -> Iterator[Image.Image]:
     # Convert a PDF in small chunks of pages at a time (e.g. 1-10, 11-20... and so on)
     exactly_one(filename=filename, file=file)
     if file is not None:
         f_bytes = convert_to_bytes(file)
-        info = pdf2image.pdfinfo_from_bytes(f_bytes)
+        info = pdf2image.pdfinfo_from_bytes(f_bytes, userpw=password)
     else:
         f_bytes = None
-        info = pdf2image.pdfinfo_from_path(filename)
+        info = pdf2image.pdfinfo_from_path(filename, userpw=password)
 
     total_pages = info["Pages"]
     for start_page in range(1, total_pages + 1, chunk_size):
@@ -407,12 +412,14 @@ def convert_pdf_to_images(
                 f_bytes,
                 first_page=start_page,
                 last_page=end_page,
+                userpw=password,
             )
         else:
             chunk_images = pdf2image.convert_from_path(
                 filename,
                 first_page=start_page,
                 last_page=end_page,
+                userpw=password,
             )
 
         for image in chunk_images: