Watermark support for PDF images (#26)

CHANGELOG.md

@@ -8,6 +8,12 @@ and this project adheres to [Semantic Versioning][].
 [keep a changelog]: https://keepachangelog.com/en/1.0.0/
 [semantic versioning]: https://semver.org/spec/v2.0.0.html
 
+## [Unreleased]
+
+### New Features
+
+-   `dso watermark` now supports files in PDF format. With this change, quarto reports using the watermark feature can
+    be rendered to PDF, too.
 
 ## v0.8.2
 

pyproject.toml

@@ -30,6 +30,7 @@ dependencies = [
   "pillow",
   "platformdirs",
   "pre-commit",
+  "pypdf",
   "pyyaml",
   "questionary",
   "rich-click",

src/dso/watermark.py

@@ -8,6 +8,7 @@
 
 import rich_click as click
 from PIL import Image, ImageDraw, ImageFont
+from pypdf import PdfReader, PdfWriter
 from svgutils import compose
 
 from dso import assets
@@ -162,10 +163,23 @@ def apply_and_save(self, input_image: Path | str, output_image: Path | str):
 class PDFWatermarker(Watermarker):
     """Add watermarks to PDF images. The watermark overlay will be a pixel graphic embedded in the svg."""
 
-    # e.g. https://www.geeksforgeeks.org/working-with-pdf-files-in-python/
+    # Inspired by https://www.geeksforgeeks.org/working-with-pdf-files-in-python/
     def apply_and_save(self, input_image: Path | str, output_image: Path | str):
         """Apply the watermark to an image and save it to the specified output file"""
-        raise NotImplementedError
+        reader = PdfReader(input_image)
+        writer = PdfWriter()
+        for page_obj in reader.pages:
+            size = (int(page_obj.mediabox.width), int(page_obj.mediabox.height))
+            watermark_overlay = self.get_watermark_overlay(size)
+            with tempfile.NamedTemporaryFile(suffix=".pdf") as tf:
+                watermark_overlay.save(tf)
+                watermark_overlay_pdf = PdfReader(tf.file).pages[0]
+                page_obj.merge_page(watermark_overlay_pdf)
+                writer.add_page(page_obj)
+
+        with open(output_image, "wb") as f:
+            writer.write(f)
+        reader.close()
 
 
 @click.command(name="watermark")

tests/test_pandocfilter.py

@@ -7,6 +7,7 @@
 
 def test_pandocfilter(quarto_stage):
     copyfile(TESTDATA / "git_logo.png", quarto_stage / "src" / "git_logo.png")
+    copyfile(TESTDATA / "git_logo.pdf", quarto_stage / "src" / "git_logo.pdf")
     copyfile(TESTDATA / "git_logo.svg", quarto_stage / "src" / "git_logo.svg")
     copyfile(TESTDATA / "git_logo.svg", quarto_stage / "src" / "git logo.svg")
     (quarto_stage / "src" / "_quarto.yml").write_text(
@@ -40,6 +41,10 @@ def test_pandocfilter(quarto_stage):
 
             ![PNG Image](git_logo.png)
 
+            a PDF Image
+
+            ![PDF Image](git_logo.pdf)
+
             and an SVG image
 
             ![SVG Image](git_logo.svg)

tests/test_watermark.py

@@ -4,7 +4,7 @@
 from click.testing import CliRunner
 from PIL import Image
 
-from dso.watermark import SVGWatermarker, Watermarker, cli
+from dso.watermark import PDFWatermarker, SVGWatermarker, Watermarker, cli
 from tests.conftest import TESTDATA
 
 
@@ -30,6 +30,18 @@ def test_add_watermark_svg(tmp_path):
     wm.apply_and_save(TESTDATA / "git_logo.svg", tmp_path / "git_logo_watermarked.svg")
 
 
+@pytest.mark.parametrize(
+    "pdf_file",
+    [
+        "git_logo.pdf",  # single page, pixel
+        "lorem_ipsum.pdf",  # multi page, vector
+    ],
+)
+def test_add_watermark_pdf(tmp_path, pdf_file):
+    wm = PDFWatermarker("test")
+    wm.apply_and_save(TESTDATA / pdf_file, tmp_path / pdf_file)
+
+
 @pytest.mark.parametrize(
     "params",
     [