当前位置：首页 > news >正文

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

news 来源：原创 2025/5/10 22:12:37

给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中。整个过程支持自定义水印的位置、角度和透明度，能够灵活地满足不同场景的需求。代码结构清晰，易于扩展和维护，适合在实际项目中使用。

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