Add with ... as ... usage

PyPDF2/_merger.py

@@ -26,7 +26,7 @@
 # POSSIBILITY OF SUCH DAMAGE.
 
 from io import BytesIO, FileIO, IOBase
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, cast
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, cast, Type
 
 from ._encryption import Encryption
 from ._page import PageObject
@@ -51,6 +51,7 @@
 )
 from .pagerange import PageRange, PageRangeSpec
 from .types import FitType, LayoutType, OutlinesType, PagemodeType, ZoomArgType
+from types import TracebackType
 
 ERR_CLOSED_WRITER = "close() was called and thus the writer cannot be used anymore"
 
@@ -75,20 +76,33 @@ class PdfMerger:
     See the functions :meth:`merge()<merge>` (or :meth:`append()<append>`)
     and :meth:`write()<write>` for usage information.
 
+    :param fileobj: Output file. Can be a filename or any kind of
+            file-like object.
     :param bool strict: Determines whether user should be warned of all
             problems and also causes some correctable problems to be fatal.
             Defaults to ``False``.
     """
 
-    def __init__(self, strict: bool = False) -> None:
+    def __init__(self, fileobj: StrByteType = "", strict: bool = False) -> None:
         self.inputs: List[Tuple[Any, PdfReader, bool]] = []
         self.pages: List[Any] = []
         self.output: Optional[PdfWriter] = PdfWriter()
         self.bookmarks: OutlinesType = []
         self.named_dests: List[Any] = []
         self.id_count = 0
+        self.fileobj = fileobj
         self.strict = strict
 
+    # There is nothing to do.
+    def __enter__(self) -> "PdfMerger":
+        return self
+
+    # Write to the fileobj and close the merger.
+    def __exit__(self, exc_type: Optional[Type[BaseException]], exc: Optional[BaseException],
+                 traceback: Optional[TracebackType]) -> None:
+        self.write(self.fileobj)
+        self.close()
+
     def merge(
         self,
         position: int,
@@ -252,10 +266,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
@@ -276,9 +286,6 @@ def write(self, fileobj: StrByteType) -> None:
         # Write the output to the file
         self.output.write(fileobj)
 
-        if my_file:
-            fileobj.close()
-
     def close(self) -> None:
         """Shut all file descriptors (input and output) and clear all memory usage."""
         self.pages = []

PyPDF2/_writer.py

@@ -47,6 +47,7 @@
     Tuple,
     Union,
     cast,
+    Type
 )
 
 from PyPDF2.errors import PdfReadWarning
@@ -59,6 +60,7 @@
     _get_max_pdf_version_header,
     b_,
     deprecate_with_replacement,
+    StrByteType
 )
 from .constants import AnnotationDictionaryAttributes
 from .constants import CatalogAttributes as CA
@@ -107,6 +109,9 @@
     ZoomArgsType,
     ZoomArgType,
 )
+from io import FileIO
+from types import TracebackType
+
 
 logger = logging.getLogger(__name__)
 
@@ -121,7 +126,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 +163,16 @@ def __init__(self) -> None:
         )
         self._root: Optional[IndirectObject] = None
         self._root_object = root
+        self.fileobj = fileobj
+
+    # Nothing to do.
+    def __enter__(self) -> "PdfWriter":
+        return self
+
+    # Write to the fileobj.
+    def __exit__(self, exc_type: Optional[Type[BaseException]], exc: Optional[BaseException],
+                 traceback: Optional[TracebackType]) -> None:
+        self.write(self.fileobj)
 
     @property
     def pdf_header(self) -> bytes:
@@ -763,7 +778,7 @@ def encrypt(
         self._encrypt = self._add_object(encrypt)
         self._encrypt_key = key
 
-    def write(self, stream: StreamType) -> None:
+    def write_stream(self, stream: StreamType) -> None:
         """
         Write the collection of pages added to this object out as a PDF file.
 
@@ -794,6 +809,23 @@ 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, fileobj: StrByteType) -> None:
+        my_file = False
+
+        if fileobj == "":
+            raise ValueError(f"Output(fileobj={fileobj}) is empty.")
+
+        if isinstance(fileobj, str):
+            fileobj = FileIO(fileobj, "wb")
+            my_file = True
+        else:
+            pass
+
+        self.write_stream(fileobj)
+
+        if my_file:
+            fileobj.close()
+
     def _write_header(self, stream: StreamType) -> List[int]:
         object_positions = []
         stream.write(self.pdf_header + b"\n")
@@ -1601,7 +1633,7 @@ def _get_page_layout(self) -> Optional[LayoutType]:
 
     def getPageLayout(self) -> Optional[LayoutType]:  # pragma: no cover
         """
-        .. deprecated:: 1.28.0
+        .. deprecated:: 1.28.0raise
 
             Use :py:attr:`page_layout` instead.
         """