utils/pdf_utils.py

import PyPDF2
import io
from typing import List

def merge_pdfs(files: List) -> bytes:
    """合并多个PDF文件"""
    merger = PyPDF2.PdfMerger()
    for file in files:
        merger.append(file)
    buffer = io.BytesIO()
    merger.write(buffer)
    merger.close()
    return buffer.getvalue()

def split_pdf(file, pages: str = None) -> List[bytes]:
    """拆分PDF文件"""
    reader = PyPDF2.PdfReader(file)
    total_pages = len(reader.pages)

    if pages:
        page_list = parse_page_range(pages, total_pages)
    else:
        page_list = list(range(total_pages))

    results = []
    for page_num in page_list:
        writer = PyPDF2.PdfWriter()
        writer.add_page(reader.pages[page_num])
        buffer = io.BytesIO()
        writer.write(buffer)
        results.append(buffer.getvalue())

    return results

def parse_page_range(pages_str: str, total_pages: int) -> List[int]:
    """解析页码范围字符串"""
    result = []
    parts = pages_str.split(',')
    for part in parts:
        part = part.strip()
        if '-' in part:
            start, end = part.split('-')
            start = int(start) - 1
            end = int(end) if end else total_pages
            result.extend(range(start, end))
        else:
            result.append(int(part) - 1)
    return result

def add_watermark(pdf_file, watermark_text: str) -> bytes:
    """添加文字水印"""
    from reportlab.lib.pagesizes import letter
    from reportlab.pdfgen import canvas

    reader = PyPDF2.PdfReader(pdf_file)
    writer = PyPDF2.PdfWriter()

    for page in reader.pages:
        # 创建水印
        watermark_buffer = io.BytesIO()
        c = canvas.Canvas(watermark_buffer, pagesize=letter)
        c.setFont("Helvetica", 50)
        c.setFillAlpha(0.3)
        c.saveState()
        c.translate(300, 400)
        c.rotate(45)
        c.drawCentredString(0, 0, watermark_text)
        c.restoreState()
        c.save()

        # 添加水印到页面
        watermark_buffer.seek(0)
        watermark_reader = PyPDF2.PdfReader(watermark_buffer)
        watermark_page = watermark_reader.pages[0]
        page.merge_page(watermark_page)
        writer.add_page(page)

    buffer = io.BytesIO()
    writer.write(buffer)
    return buffer.getvalue()

def add_password(pdf_file, user_password: str, owner_password: str = None) -> bytes:
    """添加密码保护"""
    reader = PyPDF2.PdfReader(pdf_file)
    writer = PyPDF2.PdfWriter()

    for page in reader.pages:
        writer.add_page(page)

    if not owner_password:
        owner_password = user_password

    writer.encrypt(user_password, owner_password)
    buffer = io.BytesIO()
    writer.write(buffer)
    return buffer.getvalue()