ENH: "with" support for PdfMerger and PdfWriter

CONTRIBUTORS.md

@@ -11,6 +11,7 @@ history and [GitHubs 'Contributors' feature](https://github.com/py-pdf/PyPDF2/gr
 
 ## Contributors to the pyPdf / PyPDF2 project
 
+* [JianzhengLuo](https://github.com/JianzhengLuo)
 * [Karvonen, Harry](https://github.com/Hatell/)
 * [KourFrost](https://github.com/KourFrost)
 * [Lightup1](https://github.com/Lightup1)

PyPDF2/_merger.py

@@ -27,7 +27,18 @@
 
 from io import BytesIO, FileIO, IOBase
 from pathlib import Path
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, cast
+from types import TracebackType
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+    cast,
+)
 
 from ._encryption import Encryption
 from ._page import PageObject
@@ -84,18 +95,38 @@ class PdfMerger:
     :param bool strict: Determines whether user should be warned of all
             problems and also causes some correctable problems to be fatal.
             Defaults to ``False``.
+    :param fileobj: Output file. Can be a filename or any kind of
+            file-like object.
     """
 
     @deprecate_bookmark(bookmarks="outline")
-    def __init__(self, strict: bool = False) -> None:
+    def __init__(self, strict: bool = False, fileobj: StrByteType = "") -> None:
         self.inputs: List[Tuple[Any, PdfReader, bool]] = []
         self.pages: List[Any] = []
         self.output: Optional[PdfWriter] = PdfWriter()
         self.outline: OutlineType = []
         self.named_dests: List[Any] = []
         self.id_count = 0
+        self.fileobj = fileobj
         self.strict = strict
 
+    def __enter__(self) -> "PdfMerger":
+        # There is nothing to do.
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> None:
+        """Write to the fileobj and close the merger."""
+        if self.fileobj:
+            print(f"write to {self.fileobj}")
+            self.write(self.fileobj)
+        print("nope")
+        self.close()
+
     @deprecate_bookmark(bookmark="outline_item", import_bookmarks="import_outline")
     def merge(
         self,
@@ -263,10 +294,6 @@ def write(self, fileobj: StrByteType) -> None:
         """
         if self.output is None:
             raise RuntimeError(ERR_CLOSED_WRITER)
-        my_file = False
-        if isinstance(fileobj, str):
-            fileobj = FileIO(fileobj, "wb")
-            my_file = True
 
         # Add pages to the PdfWriter
         # The commented out line below was replaced with the two lines below it
@@ -285,7 +312,7 @@ def write(self, fileobj: StrByteType) -> None:
         self._write_outline()
 
         # Write the output to the file
-        self.output.write(fileobj)
+        my_file, fileobj = self.output.write(fileobj)
 
         if my_file:
             fileobj.close()

PyPDF2/_writer.py

@@ -36,6 +36,8 @@
 import time
 import uuid
 from hashlib import md5
+from io import BufferedReader, BufferedWriter, BytesIO, FileIO
+from types import TracebackType
 from typing import (
     Any,
     Callable,
@@ -44,6 +46,7 @@
     List,
     Optional,
     Tuple,
+    Type,
     Union,
     cast,
 )
@@ -52,6 +55,7 @@
 from ._reader import PdfReader
 from ._security import _alg33, _alg34, _alg35
 from ._utils import (
+    StrByteType,
     StreamType,
     _get_max_pdf_version_header,
     b_,
@@ -121,7 +125,7 @@ class PdfWriter:
     class (typically :class:`PdfReader<PyPDF2.PdfReader>`).
     """
 
-    def __init__(self) -> None:
+    def __init__(self, fileobj: StrByteType = "") -> None:
         self._header = b"%PDF-1.3"
         self._objects: List[Optional[PdfObject]] = []  # array of indirect objects
         self._idnum_hash: Dict[bytes, IndirectObject] = {}
@@ -158,6 +162,23 @@ def __init__(self) -> None:
         )
         self._root: Optional[IndirectObject] = None
         self._root_object = root
+        self.fileobj = fileobj
+        self.with_as_usage = False
+
+    def __enter__(self) -> "PdfWriter":
+        """Store that writer is initialized by 'with'."""
+        self.with_as_usage = True
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> None:
+        """Write data to the fileobj."""
+        if self.fileobj:
+            self.write(self.fileobj)
 
     @property
     def pdf_header(self) -> bytes:
@@ -771,13 +792,7 @@ def encrypt(
         self._encrypt = self._add_object(encrypt)
         self._encrypt_key = key
 
-    def write(self, stream: StreamType) -> None:
-        """
-        Write the collection of pages added to this object out as a PDF file.
-
-        :param stream: An object to write the file to.  The object must support
-            the write method and the tell method, similar to a file object.
-        """
+    def write_stream(self, stream: StreamType) -> None:
         if hasattr(stream, "mode") and "b" not in stream.mode:
             logger_warning(
                 f"File <{stream.name}> to write to is not in binary mode. "  # type: ignore
@@ -803,6 +818,33 @@ def write(self, stream: StreamType) -> None:
         self._write_trailer(stream)
         stream.write(b_(f"\nstartxref\n{xref_location}\n%%EOF\n"))  # eof
 
+    def write(
+        self, stream: StrByteType
+    ) -> Tuple[bool, Union[FileIO, BytesIO, BufferedReader, BufferedWriter]]:
+        """
+        Write the collection of pages added to this object out as a PDF file.
+
+        :param stream: An object to write the file to.  The object can support
+            the write method and the tell method, similar to a file object, or
+            be a file path, just like the fileobj, just named it stream to keep
+            existing workflow.
+        """
+        my_file = False
+
+        if stream == "":
+            raise ValueError(f"Output(stream={stream}) is empty.")
+
+        if isinstance(stream, str):
+            stream = FileIO(stream, "wb")
+            my_file = True
+
+        self.write_stream(stream)
+
+        if self.with_as_usage:
+            stream.close()
+
+        return my_file, stream
+
     def _write_header(self, stream: StreamType) -> List[int]:
         object_positions = []
         stream.write(self.pdf_header + b"\n")

tests/test_merger.py

@@ -18,14 +18,12 @@
 sys.path.append(str(PROJECT_ROOT))
 
 
-def test_merge():
+def merger_operate(merger):
     pdf_path = RESOURCE_ROOT / "crazyones.pdf"
     outline = RESOURCE_ROOT / "pdflatex-outline.pdf"
     pdf_forms = RESOURCE_ROOT / "pdflatex-forms.pdf"
     pdf_pw = RESOURCE_ROOT / "libreoffice-writer-password.pdf"
 
-    merger = PyPDF2.PdfMerger()
-
     # string path:
     merger.append(pdf_path)
     merger.append(outline)
@@ -95,10 +93,8 @@ def test_merge():
     merger.set_page_layout("/SinglePage")
     merger.set_page_mode("/UseThumbs")
 
-    tmp_path = "dont_commit_merged.pdf"
-    merger.write(tmp_path)
-    merger.close()
 
+def check_outline(tmp_path):
     # Check if outline is correct
     reader = PyPDF2.PdfReader(tmp_path)
     assert [el.title for el in reader.outline if isinstance(el, Destination)] == [
@@ -117,7 +113,50 @@ def test_merge():
 
     # TODO: There seem to be no destinations for those links?
 
-    # Clean up
+
+tmp_path = "dont_commit_merged.pdf"
+
+
+def test_merger_operations_by_traditional_usage():
+    # Arrange
+    merger = PdfMerger()
+    merger_operate(merger)
+
+    # Act
+    merger.write(tmp_path)
+    merger.close()
+
+    # Assert
+    check_outline(tmp_path)
+
+    # cleanup
+    os.remove(tmp_path)
+
+
+def test_merger_operations_by_semi_traditional_usage():
+    with PdfMerger() as merger:
+        # Arrange
+        merger_operate(merger)
+        # Act
+        merger.write(tmp_path)
+
+    # Assert
+    assert os.path.isfile(tmp_path)
+    check_outline(tmp_path)
+
+    # cleanup
+    os.remove(tmp_path)
+
+
+def test_merger_operation_by_new_usage():
+    with PdfMerger(fileobj=tmp_path) as merger:
+        merger_operate(merger)
+
+    # Assert
+    assert os.path.isfile(tmp_path)
+    check_outline(tmp_path)
+
+    # cleanup
     os.remove(tmp_path)
 
 

tests/test_writer.py

@@ -26,20 +26,16 @@ def test_writer_clone():
     assert len(writer.pages) == 4
 
 
-def test_writer_operations():
+def writer_operate(writer):
     """
-    This test just checks if the operation throws an exception.
-
-    This should be done way more thoroughly: It should be checked if the
-    output is as expected.
+    To test the writer that initialized by each of the four usages.
     """
     pdf_path = RESOURCE_ROOT / "crazyones.pdf"
     pdf_outline_path = RESOURCE_ROOT / "pdflatex-outline.pdf"
 
     reader = PdfReader(pdf_path)
     reader_outline = PdfReader(pdf_outline_path)
 
-    writer = PdfWriter()
     page = reader.pages[0]
     with pytest.raises(PageSizeNotDefinedError) as exc:
         writer.add_blank_page()
@@ -91,17 +87,57 @@ def test_writer_operations():
 
     writer.add_attachment("foobar.gif", b"foobarcontent")
 
-    # finally, write "output" to PyPDF2-output.pdf
-    tmp_path = "dont_commit_writer.pdf"
-    with open(tmp_path, "wb") as output_stream:
-        writer.write(output_stream)
-
     # Check that every key in _idnum_hash is correct
     objects_hash = [o.hash_value() for o in writer._objects]
     for k, v in writer._idnum_hash.items():
         assert v.pdf == writer
         assert k in objects_hash, "Missing %s" % v
 
+
+tmp_path = "dont_commit_writer.pdf"
+
+
+def test_writer_operations_by_totally_traditional_usage():
+    writer = PdfWriter()
+
+    writer_operate(writer)
+
+    # finally, write "output" to PyPDF2-output.pdf
+    with open(tmp_path, "wb") as output_stream:
+        writer.write(output_stream)
+
+    # cleanup
+    os.remove(tmp_path)
+
+
+def test_writer_operations_by_semi_traditional_usage():
+    with PdfWriter() as writer:
+        writer_operate(writer)
+
+        # finally, write "output" to PyPDF2-output.pdf
+        with open(tmp_path, "wb") as output_stream:
+            writer.write(output_stream)
+
+    # cleanup
+    os.remove(tmp_path)
+
+
+def test_writer_operations_by_semi_new_traditional_usage():
+    with PdfWriter() as writer:
+        writer_operate(writer)
+
+        # finally, write "output" to PyPDF2-output.pdf
+        writer.write(tmp_path)
+
+    # cleanup
+    os.remove(tmp_path)
+
+
+def test_writer_operation_by_totally_new_usage():
+    # This includes write "output" to PyPDF2-output.pdf
+    with PdfWriter(tmp_path) as writer:
+        writer_operate(writer)
+
     # cleanup
     os.remove(tmp_path)
 
@@ -656,3 +692,13 @@ def test_colors_in_outline_item():
 
     # Cleanup
     os.remove(target)  # remove for testing
+
+
+def test_write_empty_stream():
+    reader = PdfReader(EXTERNAL_ROOT / "004-pdflatex-4-pages/pdflatex-4-pages.pdf")
+    writer = PdfWriter()
+    writer.clone_document_from_reader(reader)
+
+    with pytest.raises(ValueError) as exc:
+        writer.write("")
+    assert exc.value.args[0] == "Output(stream=) is empty."