Implement Watermarkable interface on PDFDocument

product/ERP5/Document/PDFDocument.py

@@ -28,15 +28,17 @@
 
 import tempfile, os
 
+import zope.interface
 from AccessControl import ClassSecurityInfo
 
-from Products.ERP5Type import Permissions, PropertySheet
+from Products.ERP5Type import Permissions, PropertySheet, interfaces
 from Products.ERP5.Document.Image import Image
 from Products.ERP5.Document.Document import ConversionError,\
                                             VALID_TEXT_FORMAT_LIST
 from subprocess import Popen, PIPE
 from zLOG import LOG
 import errno
+from StringIO import StringIO
 
 class PDFDocument(Image):
   """
@@ -66,6 +68,44 @@ class PDFDocument(Image):
                     , PropertySheet.Periodicity
                     )
 
+  zope.interface.implements(interfaces.IWatermarkable)
+
+  security.declareProtected(Permissions.AccessContentsInformation,
+                            'getWatermarkedData')
+  def getWatermarkedData(self, watermark_data, repeat_watermark=True,
+                         watermark_start_page=0, **kw):
+    """See interface
+
+    * watermark_data is the PDF data (as a string) to use as a watermark.
+    * If repeat_watermark is true, then the watermark will be applied on all
+      pages, otherwise it is applied only once.
+    * Watermark is applied at all pages starting watermark_start_page (this
+      index is 0 based)
+    """
+    from pyPdf import PdfFileWriter, PdfFileReader, pdf
+    self_reader = PdfFileReader(StringIO(self.getData()))
+    watermark_reader = PdfFileReader(StringIO(watermark_data))
+    watermark_page_count = watermark_reader.getNumPages()
+
+    output = PdfFileWriter()
+
+    for page_number in range(self_reader.getNumPages()):
+      self_page = self_reader.getPage(page_number)
+      watermark_page = None
+      if page_number >= watermark_start_page:
+        if repeat_watermark:
+          watermark_page = watermark_reader.getPage(
+            (page_number - watermark_start_page) % watermark_page_count)
+        elif page_number < (watermark_page_count + watermark_start_page):
+          watermark_page = watermark_reader.getPage(page_number - watermark_start_page)
+        if watermark_page is not None:
+          self_page.mergePage(watermark_page)
+      output.addPage(self_page)
+
+    outputStream = StringIO()
+    output.write(outputStream)
+    return outputStream.getvalue()
+
   # Conversion API
   def _convert(self, format, **kw):
     """

product/ERP5OOo/tests/testDms.py

@@ -1337,6 +1337,65 @@ class TestDocument(TestDocumentMixin):
     document.edit(file=upload_file)
     self.assertEquals('application/pdf', document.getContentType())
 
+  def test_PDF_watermark(self):
+    original_document = self.portal.portal_contributions.newContent(
+      file=makeFileUpload('REF-en-001.pdf'))
+    # This watermark.pdf document is a pdf with a transparent background. Such
+    # document can be created using GIMP
+    watermark_document = self.portal.portal_contributions.newContent(
+      file=makeFileUpload('watermark.pdf'))
+    watermarked_data = original_document.getWatermarkedData(
+      watermark_data=watermark_document.getData(),
+      repeat_watermark=False)
+
+    # this looks like a pdf
+    self.assertTrue(watermarked_data.startswith('%PDF-1.3'))
+
+    # and ERP5 can make a PDF Document out of it
+    watermarked_document = self.portal.document_module.newContent(
+      portal_type='PDF',
+      data=watermarked_data)
+    self.assertEqual('1', watermarked_document.getContentInformation()['Pages'])
+    self.assertNotEqual(original_document.getData(),
+      watermarked_document.getData())
+
+  def test_PDF_watermark_repeat(self):
+    # watermark a pdf, repeating the watermark
+    original_document = self.portal.portal_contributions.newContent(
+      file=makeFileUpload('Forty-Two.Pages-en-001.pdf'))
+    watermark_document = self.portal.portal_contributions.newContent(
+      file=makeFileUpload('watermark.pdf'))
+    watermarked_data = original_document.getWatermarkedData(
+      watermark_data=watermark_document.getData(),
+      repeat_watermark=True)
+
+    self.assertTrue(watermarked_data.startswith('%PDF-1.3'))
+    watermarked_document = self.portal.document_module.newContent(
+      portal_type='PDF',
+      data=watermarked_data)
+    self.assertEqual('42', watermarked_document.getContentInformation()['Pages'])
+    self.assertNotEqual(original_document.getData(),
+      watermarked_document.getData())
+
+  def test_PDF_watermark_start_page(self):
+    # watermark a pdf, starting on the second page
+    original_document = self.portal.portal_contributions.newContent(
+      file=makeFileUpload('Forty-Two.Pages-en-001.pdf'))
+    watermark_document = self.portal.portal_contributions.newContent(
+      file=makeFileUpload('watermark.pdf'))
+    watermarked_data = original_document.getWatermarkedData(
+      watermark_data=watermark_document.getData(),
+      repeat_watermark=False,
+      watermark_start_page=1) # This is 0 based.
+
+    self.assertTrue(watermarked_data.startswith('%PDF-1.3'))
+    watermarked_document = self.portal.document_module.newContent(
+      portal_type='PDF',
+      data=watermarked_data)
+    self.assertEqual('42', watermarked_document.getContentInformation()['Pages'])
+    self.assertNotEqual(original_document.getData(),
+      watermarked_document.getData())
+
   def test_Document_getStandardFilename(self):
     upload_file = makeFileUpload('metadata.pdf')
     document = self.portal.document_module.newContent(portal_type='PDF')