ENH: Add --layout parameter to up2

pdfly/cli.py

@@ -74,8 +74,10 @@ def up2(
         ),
     ],
     out: Path,
+    layout: str = typer.Option(None, "--layout", help="Optional page layout (e.g., 2x2, 3x3, 1x2, 2x1)")
 ) -> None:
-    pdfly.up2.main(pdf, out)
+    """Combine PDF pages based on the specified layout, if provided."""
+    pdfly.up2.up2_main(pdf, out, layout)
 
 
 @entry_point.command(name="cat", help=pdfly.cat.__doc__)  # type: ignore[misc]

pdfly/up2.py

@@ -29,3 +29,45 @@ def main(pdf: Path, output: Path) -> None:
     with open(output, "wb") as fp:
         writer.write(fp)
     print("done.")
+
+def up2_main(pdf: Path, output: Path, layout: str = None) -> None:
+    reader = PdfReader(str(pdf))
+    writer = PdfWriter()
+
+    if layout:
+        # Define layout configurations (columns, rows) for each grid type
+        layout_options = {
+            "2x2": (2, 2),
+            "3x3": (3, 3),
+            "1x2": (1, 2),
+            "2x1": (2, 1)
+        }
+
+        if layout not in layout_options:
+            raise ValueError(f"Unsupported layout: {layout}")
+
+        columns, rows = layout_options[layout]
+        # Adjusted to use 'mediabox' instead of 'media_box'
+        page_width = reader.pages[0].mediabox.width / columns
+        page_height = reader.pages[0].mediabox.height / rows
+
+        # Arrange pages in specified grid
+        for i in range(0, len(reader.pages), columns * rows):
+            new_page = writer.add_blank_page(width=reader.pages[0].mediabox.width,
+                                             height=reader.pages[0].mediabox.height)
+
+            for col in range(columns):
+                for row in range(rows):
+                    index = i + row * columns + col
+                    if index < len(reader.pages):
+                        page = reader.pages[index]
+                        x_position = col * page_width
+                        y_position = reader.pages[0].mediabox.height - (row + 1) * page_height
+                        new_page.merge_translated_page(page, x_position, y_position)
+    else:
+        # Default behavior: add pages without grid layout
+        for page in reader.pages:
+            writer.add_page(page)
+
+    with open(output, "wb") as f_out:
+        writer.write(f_out)

tests/test_up2_layout.py

@@ -0,0 +1,36 @@
+import os
+from pathlib import Path
+from pypdf import PdfReader
+from pdfly.up2 import up2_main
+
+# Create a simple 6-page blank PDF for testing
+def create_test_pdf(file_path, num_pages=6):
+    from pypdf import PdfWriter
+    writer = PdfWriter()
+    for _ in range(num_pages):
+        writer.add_blank_page(width=210, height=297)  # A4 size in mm
+    with open(file_path, "wb") as f:
+        writer.write(f)
+
+# File paths for the input and output PDFs
+input_pdf = Path("test_input.pdf")
+output_pdf = Path("test_output.pdf")
+
+# Create the test PDF
+create_test_pdf(input_pdf)
+
+# List of layouts to test
+layouts = ["2x2", "3x3", "1x2", "2x1"]
+for layout in layouts:
+    print(f"\nTesting layout: {layout}")
+    up2_main(input_pdf, output_pdf, layout=layout)
+
+    # Read the output PDF and print the number of pages
+    reader = PdfReader(str(output_pdf))
+    print(f"Output PDF for layout {layout} has {len(reader.pages)} pages.")
+
+# Clean up
+if output_pdf.exists():
+    os.remove(output_pdf)
+if input_pdf.exists():
+    os.remove(input_pdf)