Repository: sfneal/pdfconduit

Issue:
Add support for pdfrw actions from PDF streams
- fix issues with flatten, rotate_exact, etc methods that use pdfrw back-end
- [x] figure out how to use streams with pdfrw

### tests
- [x] can flatten pdf from stream
- [x] can rotate_exact pdf from stream
- [x] [add more as found]

Discussion:

Patch:
diff --git a/pdfconduit/_version.py b/pdfconduit/_version.py
index 7039708..72aa758 100644
--- a/pdfconduit/_version.py
+++ b/pdfconduit/_version.py
@@ -1,1 +1,1 @@
-__version__ = "4.1.0"
+__version__ = "4.1.1"
diff --git a/pdfconduit/internals/base.py b/pdfconduit/internals/base.py
index 3bb457b..22f5729 100644
--- a/pdfconduit/internals/base.py
+++ b/pdfconduit/internals/base.py
@@ -21,6 +21,7 @@ class BaseConduit(ABC):
     _closed: bool = False
 
     _pdf_file = None
+    _stream: BytesIO = None
     _reader: PdfReader
     _writer: PdfWriter
 
@@ -55,7 +56,8 @@ class BaseConduit(ABC):
         return self
 
     def _read_stream(self, stream: BytesIO) -> Self:
-        self._reader = PdfReader(stream)
+        self._stream = stream
+        self._reader = PdfReader(self._stream)
         self._writer: PdfWriter = PdfWriter(clone_from=self._reader)
         return self
 
@@ -80,6 +82,8 @@ class BaseConduit(ABC):
         self._reader.close()
         if self._pdf_file is not None:
             self._pdf_file.close()
+        if self._stream is not None:
+            self._stream.close()
 
         # Confirm output path is set
         if self.output is None:
diff --git a/pdfconduit/pdfconduit.py b/pdfconduit/pdfconduit.py
index 8253723..9606a1b 100644
--- a/pdfconduit/pdfconduit.py
+++ b/pdfconduit/pdfconduit.py
@@ -1,3 +1,5 @@
+from tempfile import TemporaryFile, NamedTemporaryFile
+
 from pypdf import PdfWriter
 
 from pypdf import PdfWriter
@@ -5,7 +7,7 @@ from pypdf import PdfWriter
 from pdfconduit.convert import Flatten
 from pdfconduit.internals import BaseConduit
 from pdfconduit.settings import Compression, ImageQualityRange, Encryption
-from pdfconduit.transform import Merge
+from pdfconduit.transform import Merge2
 from pdfconduit.transform import Rotate, Upscale
 from pdfconduit.utils import Info
 from pdfconduit.utils.typing import Optional, Tuple, Self, Annotated
@@ -33,9 +35,8 @@ class Pdfconduit(BaseConduit):
 
     def merge_fast(self, pdfs: list) -> Self:
         self._set_default_output("merged")
-        self._path = (
-            Merge([self._path] + pdfs, output_dir=self._output_dir).use_pdfrw().merge()
-        )
+        pdf_objects = [self._stream if self._stream is not None else self._path] + pdfs
+        self._path = Merge2(pdf_objects, output=self.output).use_pdfrw().merge()
         return self._open_and_read()
 
     def rotate(self, degrees: int) -> Self:
@@ -48,7 +49,16 @@ class Pdfconduit(BaseConduit):
         if degrees % 90 == 0:
             return self.rotate(degrees)
 
-        self._path = Rotate(self._path, degrees).use_pdfrw().rotate()
+        self._set_default_output("rotated")
+        self._path = (
+            Rotate(
+                self._stream if self._stream is not None else self._path,
+                degrees,
+                output=self.output,
+            )
+            .use_pdfrw()
+            .rotate()
+        )
         return self._open_and_read()
 
     def slice(self, start: int, end: int) -> Self:
@@ -89,11 +99,22 @@ class Pdfconduit(BaseConduit):
     def flatten(self) -> Self:
         # todo: re-write Flatten & other convert classes
         # todo: fix issue with flattened pdf output path
+        if not self._path and self._stream:
+            temp = NamedTemporaryFile(suffix=".pdf")
+            temp.write(self._stream.getvalue())
+            path = temp.name
+        else:
+            temp = None
+            path = self._path
+
         if not self._closed:
             self.write()
-        self._path = Flatten(
-            self._path, suffix="flattened", tempdir=self._output_dir
-        ).save()
+
+        self._path = Flatten(path, suffix="flattened", tempdir=self._output_dir).save()
+
+        if temp is not None:
+            temp.close()
+
         return self._open_and_read()
 
     def minify(self) -> Self:
diff --git a/pdfconduit/transform/__init__.py b/pdfconduit/transform/__init__.py
index 68cea6a..5e4614f 100644
--- a/pdfconduit/transform/__init__.py
+++ b/pdfconduit/transform/__init__.py
@@ -1,5 +1,6 @@
 from pdfconduit.transform.merge import Merge
+from pdfconduit.transform.merge2 import Merge2
 from pdfconduit.transform.rotate import Rotate
 from pdfconduit.transform.upscale import Upscale
 
-__all__ = [Merge, Rotate, Upscale]
+__all__ = [Merge, Merge2, Rotate, Upscale]
diff --git a/pdfconduit/transform/merge2.py b/pdfconduit/transform/merge2.py
new file mode 100644
index 0000000..129f495
--- /dev/null
+++ b/pdfconduit/transform/merge2.py
@@ -0,0 +1,51 @@
+# Merge PDF documents
+from io import BytesIO
+from typing import Union, Iterable
+
+from pdfrw import (
+    PdfReader as PdfrwReader,
+    PdfWriter as PdfrwWriter,
+    IndirectPdfDict as PdfrwIndirectPdfDict,
+)
+from pypdf import PdfWriter as PyPdfWriter
+
+from pdfconduit.utils.driver import PdfDriver
+
+
+class Merge2(PdfDriver):
+    def __init__(self, pdfs: Iterable[Union[str, BytesIO]], output: str):
+        self._pdfs = pdfs
+        self._output = output
+
+    def merge(self) -> str:
+        """Merge list of PDF files to a single PDF file."""
+        return self.execute()
+
+    def pdfrw(self):
+        writer = PdfrwWriter()
+
+        for pdf_object in self._pdfs:
+            if isinstance(pdf_object, BytesIO):
+                reader = PdfrwReader(fdata=pdf_object.getvalue())
+            else:
+                reader = PdfrwReader(fname=pdf_object)
+            writer.addpages(reader.pages)
+
+        writer.trailer.Info = PdfrwIndirectPdfDict(
+            Author="Stephen Neal",
+            Creator="pdfconduit",
+            Producer="pdfconduit",
+        )
+        writer.write(self._output)
+        return self._output
+
+    def pypdf(self):
+        merger = PyPdfWriter()
+
+        for pdf in self._pdfs:
+            merger.append(pdf)
+
+        merger.write(self._output)
+        merger.close()
+
+        return self._output
diff --git a/pdfconduit/transform/rotate.py b/pdfconduit/transform/rotate.py
index 784de03..e88ccd1 100644
--- a/pdfconduit/transform/rotate.py
+++ b/pdfconduit/transform/rotate.py
@@ -1,7 +1,8 @@
 # Rotate a pdf file
 import os
+from io import BytesIO
 from tempfile import NamedTemporaryFile
-from typing import Optional
+from typing import Optional, Union
 
 from pdfrw import (
     PdfReader as PdfrwReader,
@@ -17,25 +18,30 @@ class Rotate(PdfDriver):
 
     def __init__(
         self,
-        file_name: str,
+        pdf: Union[str, BytesIO],
         rotation: int,
         suffix: str = "rotated",
         tempdir: Optional[str] = None,
+        output: Optional[str] = None,
     ):
-        self.file_name = file_name
+        self.pdf_object = pdf
         self.rotation = rotation
         self.suffix = suffix
-        self.tempdir = tempdir
-
-        if tempdir:
-            with NamedTemporaryFile(suffix=".pdf", dir=tempdir, delete=False) as temp:
-                self.outfn = temp.name
-        elif suffix:
-            self.outfn = os.path.join(
-                os.path.dirname(file_name), add_suffix(file_name, suffix)
-            )
+
+        if output:
+            self.outfn = output
         else:
-            self.outfn = NamedTemporaryFile(suffix=".pdf").name
+            self.tempdir = tempdir
+
+            if tempdir:
+                with NamedTemporaryFile(
+                    suffix=".pdf", dir=tempdir, delete=False
+                ) as temp:
+                    self.outfn = temp.name
+            elif suffix:
+                self.outfn = os.path.join(os.path.dirname(pdf), add_suffix(pdf, suffix))
+            else:
+                self.outfn = NamedTemporaryFile(suffix=".pdf").name
 
     def __str__(self) -> str:
         return self.file
@@ -48,7 +54,11 @@ class Rotate(PdfDriver):
         return str(self.outfn)
 
     def pdfrw(self) -> str:
-        trailer = PdfrwReader(self.file_name)
+        if isinstance(self.pdf_object, BytesIO):
+            trailer = PdfrwReader(fdata=self.pdf_object.getvalue())
+        else:
+            trailer = PdfrwReader(fname=self.pdf_object)
+
         pages = trailer.pages
 
         ranges = [[1, len(pages)]]
@@ -66,7 +76,7 @@ class Rotate(PdfDriver):
         return self.outfn
 
     def pypdf(self) -> str:
-        reader = PypdfReader(self.file_name)
+        reader = PypdfReader(self.pdf_object)
         writer = PypdfWriter()
 
         for page_num in range(1, reader.get_num_pages()):


Test patch:
diff --git a/tests/__init__.py b/tests/__init__.py
index 5b8aff4..c144c1f 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -4,8 +4,10 @@ import os
 import shutil
 import unittest
 from decimal import Decimal
+from io import BytesIO
 from tempfile import TemporaryDirectory
 from time import time
+from typing import Optional
 
 from pdfconduit import Pdfconduit, Info
 from pdfconduit.utils.typing import Iterable
@@ -145,6 +147,10 @@ class PdfconduitTestCase(unittest.TestCase):
     def assertFileSizeDecreased(self, original: str, modified: str):
         self.assertLess(os.path.getsize(modified), os.path.getsize(original))
 
+    def _get_pdf_byte_stream(self, path: Optional[str] = None) -> BytesIO:
+        with open(path if path else self.pdf_path, "rb") as fh:
+            return BytesIO(fh.read())
+
 
 class Timer:
     def __init__(self, decimal_places=2):
diff --git a/tests/test_flatten.py b/tests/test_flatten.py
index 4d60890..f5fb1df 100644
--- a/tests/test_flatten.py
+++ b/tests/test_flatten.py
@@ -33,3 +33,14 @@ class TestFlatten(PdfconduitTestCase):
         self.assertPdfExists(self.conduit.output)
         self.assertPdfPagesEqual(pdf_path, self.conduit.output)
         # todo: improve assertions
+
+    @parameterized.expand(flatten_params, name_func=flatten_name_func)
+    def test_flatten_from_stram(self, pdf_path: str):
+        self.conduit = Pdfconduit(
+            self._get_pdf_byte_stream(pdf_path)
+        ).set_output_directory(self.temp.name)
+        self.conduit.flatten().write()
+
+        self.assertPdfExists(self.conduit.output)
+        self.assertPdfPagesEqual(pdf_path, self.conduit.output)
+        # todo: improve assertions
diff --git a/tests/test_merge.py b/tests/test_merge.py
index a8ada3b..fd3b96b 100644
--- a/tests/test_merge.py
+++ b/tests/test_merge.py
@@ -68,3 +68,19 @@ class TestMerge(PdfconduitTestCase):
 
         self.assertPdfExists(self.conduit.output)
         self.assertCorrectNumPages(main_pdf, pdfs_to_merge, self.conduit.info.pages)
+
+    @parameterized.expand(merge_params, name_func=merge_name_func)
+    def test_can_merge_pdfs_fast_from_stream(
+        self, main_pdf: str, pdfs_to_merge: List[str]
+    ):
+        self.pdf_path = main_pdf
+        self.conduit = Pdfconduit(
+            self._get_pdf_byte_stream(self.pdf_path)
+        ).set_output_directory(self.temp.name)
+
+        self.conduit.merge_fast(pdfs_to_merge)
+
+        self.conduit.write()
+
+        self.assertPdfExists(self.conduit.output)
+        self.assertCorrectNumPages(main_pdf, pdfs_to_merge, self.conduit.info.pages)
diff --git a/tests/test_pdfconduit.py b/tests/test_pdfconduit.py
index d3335b9..601cbd8 100644
--- a/tests/test_pdfconduit.py
+++ b/tests/test_pdfconduit.py
@@ -1,7 +1,6 @@
 import os
 import warnings
-from io import BufferedReader, BytesIO
-from typing import Optional
+from io import BufferedReader
 
 from pypdf import PdfReader, PdfWriter
 
@@ -114,7 +113,3 @@ class TestUsage(PdfconduitTestCase):
 
         self.conduit.cleanup()
         self.assertPdfDoesntExists(self.conduit.output)
-
-    def _get_pdf_byte_stream(self, path: Optional[str] = None) -> BytesIO:
-        with open(path if path else self.pdf_path, "rb") as fh:
-            return BytesIO(fh.read())
diff --git a/tests/test_rotate.py b/tests/test_rotate.py
index 8c6a294..95f6b7a 100644
--- a/tests/test_rotate.py
+++ b/tests/test_rotate.py
@@ -3,7 +3,7 @@ from typing import List
 
 from parameterized import parameterized
 
-from pdfconduit import Info
+from pdfconduit import Info, Pdfconduit
 from tests import PdfconduitTestCase
 
 
@@ -56,3 +56,15 @@ class TestRotate(PdfconduitTestCase):
         self.assertTrue(
             "Rotation angle must be a multiple of 90" in str(context.exception)
         )
+
+    @parameterized.expand(rotate_exact_params, name_func=rotate_name_func)
+    def test_can_rotate_exact_from_stream(self, rotation: int):
+        self.conduit = Pdfconduit(
+            self._get_pdf_byte_stream(self.pdf_path)
+        ).set_output_directory(self.temp.name)
+        self.conduit.rotate_exact(rotation).set_output_suffix(
+            "rotated_{}".format(rotation)
+        ).write()
+
+        self.assertPdfExists(self.conduit.output)
+        self.assertPdfRotation(self.conduit, rotation)