Repository: py-pdf/pypdf
Pull request: ENH: "with" support for PdfMerger and PdfWriter
Closes #1108
Closes #1117

Co-authored-by: JianzhengLuo <jianzheng.luo.china@gmail.com>

Problem statement:
Add `with ... as ...` Usage to `PdfMerger()`
## Explanation

As I need to call `write()` and `close()` in `PdfMerger()`, why not add  `__enter__` and `__exit__` method to it so I can use it more elegantly.

## Code Example

```python3
from PyPDF2 import PdfMerger

with PdfMerger(strict=False) as merger:
    for name in ["01.pdf", "02.pdf", "03.pdf", ...]:
        merger.append(name)
    merger.write("./output.pdf")
```

Hints from issue discussion:
I like the idea! Do you want to create a pr? 
I'd like so, but as a middle school student in China, I am still busy even in summer vacation, but I will try my BEST!
I noticed that `write()` is also necessary, and it has a parameter called `fileobj`, so I decided to add `fileobj` to the class initialization, now the code looks like this:

```python3
from PyPDF2 import PdfMerger

with PdfMerger(strict=False, fileobj="./output.pdf") as merger:
    for name in ["01.pdf", "02.pdf", "03.pdf", ...]:
        merger.append(name)
```

`PdfWriter()` also need to call `write()`, but it required `stream`, and the `stream` need to be writable, that's not as convenient as `fileboj` I thought, so I made it just like `PdfMerger()`----set `fileobj` when initializing.

It will be like:

```python3
from PyPDF2 import PdfWriter

with PdfWriter(fileobj="./output.pdf") as writer:
    ...  # Do what you like to do.
```

Commit message: Add JianzhengLuo to contributors list

Patch to review:
diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md
index 28d3d9b8d..23f1ee1a7 100644
--- a/CONTRIBUTORS.md
+++ b/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)
diff --git a/PyPDF2/_merger.py b/PyPDF2/_merger.py
index db6bf0480..930ca9827 100644
--- a/PyPDF2/_merger.py
+++ b/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()
diff --git a/PyPDF2/_writer.py b/PyPDF2/_writer.py
index 6f54df2ff..e5a422084 100644
--- a/PyPDF2/_writer.py
+++ b/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")
diff --git a/tests/test_merger.py b/tests/test_merger.py
index f4cf78179..9480fdeec 100644
--- a/tests/test_merger.py
+++ b/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)
 
 
diff --git a/tests/test_writer.py b/tests/test_writer.py
index 048bd4ed7..668fd9eff 100644
--- a/tests/test_writer.py
+++ b/tests/test_writer.py
@@ -26,12 +26,9 @@ 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"
@@ -39,7 +36,6 @@ def test_writer_operations():
     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."


Code review comments:
diff_hunk: @@ -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:
path: PyPDF2/_merger.py
text: @user1:
Should `fileobj` be allowed to be path-like following from #1190?

Additionally, why not make the type `Optional[StrByteType]` and just have it default to None instead of `""`?

@author:
Good idea to allow Path objects :+1: 

What is the advantage of using `None`?

diff_hunk: @@ -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")
path: PyPDF2/_merger.py
text: @user1:
Remove these prints?

@author:
:sweat_smile: 

Should we add https://pypi.org/project/flake8-print/ ?

diff_hunk: @@ -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()
path: PyPDF2/_merger.py
text: @user1:
My preference would be to avoid overwriting the `fileobj` passed into the surrounding `write` function with got returned from `self.output.write`:
```suggestion
        my_file, ret_fileobj = self.output.write(fileobj)

        if my_file:
            ret_fileobj.close()
```

It makes it more clear on what's getting operated on imo.

diff_hunk: @@ -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)
path: tests/test_writer.py
text: @user1:
as a point of interest, if something does error within this test, do we want this file to be left behind under some circumstance?

@author:
For the writer, I don't think so. Is there a better way to do the cleanup?

For some other tests I sometimes look at the output (e.g. for the annotations). I want to see that it looks right; so during development I comment the cleanup out.

@user1:
I'd suggest using a [`tempfile.TemporaryDirectory`](https://docs.python.org/3/library/tempfile.html#tempfile.TemporaryDirectory) class. There's also [`tempfile.TemporaryFile`](https://docs.python.org/3/library/tempfile.html#tempfile.TemporaryFile) if you want to act directly on a file-like object, but it has caveats around opening it multiple times, so may not be suitable, vs `TemporaryDirectory` which just gives us a new empty directory we can create whatever files into, with the expectation that `PdfWriter` or whatever class will be opening/closing the files.

Something like:
```python
import tempfile
with tempfile.TemporaryDirectory() as tmp:
    tmp_path = os.path.join(tmp, 'file.pdf')
    writer = PdfWriter()

    writer_operate(writer)

    # finally, write "output" to PyPDF2-output.pdf
    with open(tmp_path, "wb") as output_stream:
        writer.write(output_stream)
```

or for `TemporaryFile`:

```python
import tempfile
with tempfile.TemporaryFile() as tmp_file:
    writer = PdfWriter()

    writer_operate(writer)

    # finally, write "output" to PyPDF2-output.pdf
    writer.write(tmp_file)
```

where then on leaving the context manager, `TemporaryDirectory()` will automatically remove the directory. I've not tested the above, so they may need some slight tweaks.

@author:
Thank you for the suggestion :+1: 

It feels to me that this makes the tests more noisy, but [the `tmp_path` fixture](https://docs.pytest.org/en/7.1.x/how-to/tmp_path.html#the-tmp-path-fixture) seems to follow the same idea :-) 

@author:
I've added the fixture - thank you :hugs: