【DuodooBMS】给PDF附件加“受控”水印的完整Python实现

给PDF附件加“受控”水印的完整Python实现

功能需求

在实际工作中，许多文件需要添加水印以标识其状态，例如“受控”“机密”等。对于PDF文件，添加水印不仅可以增强文件的可识别性，还可以防止未经授权的使用。本代码的功能需求是：

1. 修复PDF文件：在添加水印之前，确保PDF文件是完整且可读的，避免因文件损坏导致操作失败。

2. 添加水印：在PDF的每一页上添加指定的水印图像或文字，水印可以设置位置、角度和透明度。

3. 保存输出：将添加水印后的PDF文件保存到指定路径，并返回其二进制数据以便后续处理。

实现过程

1. 修复PDF文件

   • 使用PyMuPDF库打开PDF文件，并尝试修复。如果文件损坏，PyMuPDF可以尝试修复并保存为一个新的二进制流。

   • 如果修复失败，则直接返回原始的PDF二进制数据。

   Python复制

   def repair_pdf(self, input_pdf_binary):try:# 使用 PyMuPDF 打开并修复 PDFdoc = fitz.open(stream=input_pdf_binary, filetype="pdf")repaired_pdf_binary = BytesIO()doc.save(repaired_pdf_binary)doc.close()repaired_pdf_binary.seek(0)return repaired_pdf_binary.read()except Exception as e:print(f"Error repairing PDF: {e}")return input_pdf_binary

2. 添加水印

   • 使用reportlab库创建一个临时的PDF文件作为水印。水印可以是图像或文字，支持设置位置、角度和透明度。

   • 使用PyPDF2库将水印PDF与原始PDF合并。通过merge_page方法，将水印添加到每一页。

   Python复制

   def add_watermark(self, input_pdf_binary, output_pdf, watermark_image, x_position=30, y_position=50, opacity=1):# 尝试修复 PDFrepaired_pdf_binary = self.repair_pdf(input_pdf_binary)input_pdf_obj = PdfReader(BytesIO(repaired_pdf_binary))  # 从二进制数据中读取 PDFoutput_pdf_obj = PdfWriter()output_buffer = BytesIO()# 创建一个临时的 PDF 作为水印page_width, page_height = A4[1], A4[0]c = canvas.Canvas(output_buffer, pagesize=(page_width, page_height))try:c.setFillColor(colors.white)  # 将背景设置为白色c.setFillColor(colors.red)  # 设置字体颜色为红色c.setFont("Helvetica", 12)  # 设置字体和字体大小c.setFillAlpha(opacity)  # 设置透明度c.setStrokeColor(colors.transparent)  # 设置笔触颜色为透明img = ImageReader(watermark_image)if x_position is not None and y_position is not None:c.saveState()c.translate(x_position, y_position)c.rotate(20)c.drawImage(img, 0, 0, width=60, height=25)c.restoreState()else:x = (page_width - 60) / 2y = (page_height - 25) / 2c.saveState()c.translate(x, y)c.rotate(20)c.drawImage(img, 0, 0, width=60, height=25)c.restoreState()except Exception as e:raise ValueError(f"Error drawing image: {e}")c.showPage()c.save()# 将水印 PDF 与原始 PDF 合并watermark_pdf = PdfReader(output_buffer)for page in input_pdf_obj.pages:page.merge_page(watermark_pdf.pages[0])output_pdf_obj.add_page(page)# 保存输出 PDFfinal_output_buffer = BytesIO()output_pdf_obj.write(final_output_buffer)binary_data = final_output_buffer.getvalue()with open(output_pdf, 'wb') as f:output_pdf_obj.write(f)return binary_data

3. 调用示例

   • 准备一个PDF文件和一个水印图像文件。

   • 调用add_watermark方法，指定输入PDF、输出路径、水印图像路径等参数。

   Python复制

   if __name__ == "__main__":from io import BytesIOfrom PyPDF2 import PdfReader, PdfWriterfrom reportlab.pdfgen import canvasfrom reportlab.lib.pagesizes import A4from reportlab.lib import colorsfrom reportlab.lib.utils import ImageReaderimport fitzclass WatermarkPDF:def repair_pdf(self, input_pdf_binary):try:doc = fitz.open(stream=input_pdf_binary, filetype="pdf")repaired_pdf_binary = BytesIO()doc.save(repaired_pdf_binary)doc.close()repaired_pdf_binary.seek(0)return repaired_pdf_binary.read()except Exception as e:print(f"Error repairing PDF: {e}")return input_pdf_binarydef add_watermark(self, input_pdf_binary, output_pdf, watermark_image, x_position=30, y_position=50, opacity=1):repaired_pdf_binary = self.repair_pdf(input_pdf_binary)input_pdf_obj = PdfReader(BytesIO(repaired_pdf_binary))output_pdf_obj = PdfWriter()output_buffer = BytesIO()page_width, page_height = A4[1], A4[0]c = canvas.Canvas(output_buffer, pagesize=(page_width, page_height))try:c.setFillColor(colors.white)c.setFillColor(colors.red)c.setFont("Helvetica", 12)c.setFillAlpha(opacity)c.setStrokeColor(colors.transparent)img = ImageReader(watermark_image)if x_position is not None and y_position is not None:c.saveState()c.translate(x_position, y_position)c.rotate(20)c.drawImage(img, 0, 0, width=60, height=25)c.restoreState()else:x = (page_width - 60) / 2y = (page_height - 25) / 2c.saveState()c.translate(x, y)c.rotate(20)c.drawImage(img, 0, 0, width=60, height=25)c.restoreState()except Exception as e:raise ValueError(f"Error drawing image: {e}")c.showPage()c.save()watermark_pdf = PdfReader(output_buffer)for page in input_pdf_obj.pages:page.merge_page(watermark_pdf.pages[0])output_pdf_obj.add_page(page)final_output_buffer = BytesIO()output_pdf_obj.write(final_output_buffer)binary_data = final_output_buffer.getvalue()with open(output_pdf, 'wb') as f:output_pdf_obj.write(f)return binary_data# 示例调用watermark_pdf = WatermarkPDF()with open("example.pdf", "rb") as f:input_pdf_binary = f.read()watermark_image = "watermark.png"output_pdf = "output_with_watermark.pdf"watermark_pdf.add_watermark(input_pdf_binary, output_pdf, watermark_image)

实现总结

本代码通过PyMuPDF修复PDF文件，使用reportlab创建水印PDF，并通过PyPDF2将水印合并到原始PDF中。整个过程支持自定义水印的位置、角度和透明度，能够灵活地满足不同场景的需求。代码结构清晰，易于扩展和维护，适合在实际项目中使用。

 

让转型不迷航——邹工转型手札