给PDF附件加“受控”水印的完整Python实现
功能需求
在实际工作中，许多文件需要添加水印以标识其状态，例如“受控”“机密”等。对于PDF文件，添加水印不仅可以增强文件的可识别性，还可以防止未经授权的使用。本代码的功能需求是：
-
修复PDF文件：在添加水印之前，确保PDF文件是完整且可读的，避免因文件损坏导致操作失败。
-
添加水印：在PDF的每一页上添加指定的水印图像或文字，水印可以设置位置、角度和透明度。
-
保存输出：将添加水印后的PDF文件保存到指定路径，并返回其二进制数据以便后续处理。
实现过程
-
修复PDF文件
-
使用
PyMuPDF库打开PDF文件，并尝试修复。如果文件损坏，PyMuPDF可以尝试修复并保存为一个新的二进制流。 -
如果修复失败，则直接返回原始的PDF二进制数据。
Python复制
def repair_pdf(self, input_pdf_binary): try: # 使用 PyMuPDF 打开并修复 PDF 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_binary -
-
添加水印
-
使用
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): # 尝试修复 PDF repaired_pdf_binary = self.repair_pdf(input_pdf_binary) input_pdf_obj = PdfReader(BytesIO(repaired_pdf_binary)) # 从二进制数据中读取 PDF output_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) / 2 y = (page_height - 25) / 2 c.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) # 保存输出 PDF 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 -
-
调用示例
-
准备一个PDF文件和一个水印图像文件。
-
调用
add_watermark方法，指定输入PDF、输出路径、水印图像路径等参数。
Python复制
if __name__ == "__main__": from io import BytesIO from PyPDF2 import PdfReader, PdfWriter from reportlab.pdfgen import canvas from reportlab.lib.pagesizes import A4 from reportlab.lib import colors from reportlab.lib.utils import ImageReader import fitz class 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_binary def 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) / 2 y = (page_height - 25) / 2 c.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中。整个过程支持自定义水印的位置、角度和透明度，能够灵活地满足不同场景的需求。代码结构清晰，易于扩展和维护，适合在实际项目中使用。
让转型不迷航——邹工转型手札
