ENH: Complete attachments functions (get/add) (#1611)

Fixes #1047
Fixes #527
Fixes #169

Author: pubpub-zz

pypdf/_reader.py

@@ -2130,6 +2130,79 @@ def rename_form_topname(self, name: str) -> Optional[DictionaryObject]:
         interim[NameObject("/T")] = TextStringObject(name)
         return interim
 
+    def _list_attachments(self) -> List[str]:
+        """
+        Retrieves the list of filenames of file attachments.
+
+        Returns:
+            list of filenames
+        """
+        catalog = cast(DictionaryObject, self.trailer["/Root"])
+        # From the catalog get the embedded file names
+        try:
+            filenames = cast(
+                ArrayObject,
+                cast(
+                    DictionaryObject,
+                    cast(DictionaryObject, catalog["/Names"])["/EmbeddedFiles"],
+                )["/Names"],
+            )
+        except KeyError:
+            return []
+        attachments_names = []
+        # Loop through attachments
+        for f in filenames:
+            if isinstance(f, str):
+                attachments_names.append(f)
+        return attachments_names
+
+    def _get_attachments(
+        self, filename: Optional[str] = None
+    ) -> Dict[str, Union[bytes, List[bytes]]]:
+        """
+        Retrieves all or selected file attachments of the PDF as a dictionary of file names
+        and the file data as a bytestring.
+
+        Args:
+            filename: If filename is None, then a dictionary of all attachments
+                will be returned, where the key is the filename and the value
+                is the content. Otherwise, a dictionary with just a single key
+                - the filename - and its content will be returned.
+
+        Returns:
+            dictionary of filename -> Union[bytestring or List[ByteString]]
+            if the filename exists multiple times a List of the different version will be provided
+        """
+        catalog = cast(DictionaryObject, self.trailer["/Root"])
+        # From the catalog get the embedded file names
+        try:
+            filenames = cast(
+                ArrayObject,
+                cast(
+                    DictionaryObject,
+                    cast(DictionaryObject, catalog["/Names"])["/EmbeddedFiles"],
+                )["/Names"],
+            )
+        except KeyError:
+            return {}
+        attachments: Dict[str, Union[bytes, List[bytes]]] = {}
+        # Loop through attachments
+        for i in range(len(filenames)):
+            f = filenames[i]
+            if isinstance(f, str):
+                if filename is not None and f != filename:
+                    continue
+                name = f
+                f_dict = filenames[i + 1].get_object()
+                f_data = f_dict["/EF"]["/F"].get_data()
+                if name in attachments:
+                    if not isinstance(attachments[name], list):
+                        attachments[name] = [attachments[name]]  # type:ignore
+                    attachments[name].append(f_data)  # type:ignore
+                else:
+                    attachments[name] = f_data
+        return attachments
+
 
 class PdfFileReader(PdfReader):  # deprecated
     def __init__(self, *args: Any, **kwargs: Any) -> None:

pypdf/_writer.py

@@ -660,7 +660,7 @@ def add_attachment(self, filename: str, data: Union[str, bytes]) -> None:
         # >>
 
         ef_entry = DictionaryObject()
-        ef_entry.update({NameObject("/F"): file_entry})
+        ef_entry.update({NameObject("/F"): self._add_object(file_entry)})
 
         filespec = DictionaryObject()
         filespec.update(
@@ -685,21 +685,25 @@ def add_attachment(self, filename: str, data: Union[str, bytes]) -> None:
         # >>
         # endobj
 
-        embedded_files_names_dictionary = DictionaryObject()
-        embedded_files_names_dictionary.update(
-            {
-                NameObject(CA.NAMES): ArrayObject(
-                    [create_string_object(filename), filespec]
-                )
-            }
-        )
-
-        embedded_files_dictionary = DictionaryObject()
-        embedded_files_dictionary.update(
-            {NameObject("/EmbeddedFiles"): embedded_files_names_dictionary}
+        if CA.NAMES not in self._root_object:
+            self._root_object[NameObject(CA.NAMES)] = self._add_object(
+                DictionaryObject()
+            )
+        if "/EmbeddedFiles" not in cast(DictionaryObject, self._root_object[CA.NAMES]):
+            embedded_files_names_dictionary = DictionaryObject(
+                {NameObject(CA.NAMES): ArrayObject()}
+            )
+            cast(DictionaryObject, self._root_object[CA.NAMES])[
+                NameObject("/EmbeddedFiles")
+            ] = self._add_object(embedded_files_names_dictionary)
+        else:
+            embedded_files_names_dictionary = cast(
+                DictionaryObject,
+                cast(DictionaryObject, self._root_object[CA.NAMES])["/EmbeddedFiles"],
+            )
+        cast(ArrayObject, embedded_files_names_dictionary[CA.NAMES]).extend(
+            [create_string_object(filename), filespec]
         )
-        # Update the root
-        self._root_object.update({NameObject(CA.NAMES): embedded_files_dictionary})
 
     def addAttachment(self, fname: str, fdata: Union[str, bytes]) -> None:  # deprecated
         """

tests/test_writer.py

@@ -1194,6 +1194,34 @@ def test_iss1601():
     )
 
 
+def test_attachments():
+    writer = PdfWriter()
+    writer.add_blank_page(100, 100)
+    b = BytesIO()
+    writer.write(b)
+    b.seek(0)
+    reader = PdfReader(b)
+    b = None
+    assert reader._list_attachments() == []
+    assert reader._get_attachments() == {}
+    writer.add_attachment("foobar.txt", b"foobarcontent")
+    writer.add_attachment("foobar2.txt", b"foobarcontent2")
+    writer.add_attachment("foobar2.txt", b"2nd_foobarcontent")
+
+    b = BytesIO()
+    writer.write(b)
+    b.seek(0)
+    reader = PdfReader(b)
+    b = None
+    assert reader._list_attachments() == ["foobar.txt", "foobar2.txt", "foobar2.txt"]
+    att = reader._get_attachments()
+    assert len(att) == 2
+    assert att["foobar.txt"] == b"foobarcontent"
+    att = reader._get_attachments("foobar2.txt")
+    assert len(att) == 1
+    assert att["foobar2.txt"] == [b"foobarcontent2", b"2nd_foobarcontent"]
+
+
 @pytest.mark.external
 def test_iss1614():
     # test of an annotation(link) directly stored in the /Annots in the page