File size: 2,264 Bytes
734b5b4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | from __future__ import annotations
from pathlib import Path
import fitz
from pypdf import PdfReader, PdfWriter
from src.core.models import PDFOperation
class OfflinePDFEditingAgent:
name = "editing_agent"
def run(self, input_pdf: Path, output_pdf: Path, operations: list[PDFOperation]) -> Path:
doc = fitz.open(input_pdf)
page_ops = [op for op in operations if op.type in ("insert_text","whiteout_region","rotate_page")]
delete_ops = [op for op in operations if op.type == "delete_page"]
reorder_ops = [op for op in operations if op.type == "reorder_pages"]
for op in page_ops:
if op.page is None or op.page >= len(doc): continue
page = doc[op.page]
if op.type == "insert_text" and op.position and op.text:
color = tuple(op.color) if op.color else (0, 0, 0)
page.insert_text(op.position, op.text, fontsize=12, color=color, overlay=True)
elif op.type == "whiteout_region" and op.rect:
page.draw_rect(fitz.Rect(op.rect), color=(1,1,1), fill=(1,1,1), overlay=True)
if op.text:
x0, y0 = op.rect[0], op.rect[1]
page.insert_text((x0+4, y0+14), op.text, fontsize=12, color=(0,0,0), overlay=True)
elif op.type == "rotate_page" and op.rotation is not None:
page.set_rotation(op.rotation)
stage = output_pdf.with_name(output_pdf.stem + "-pre.pdf")
stage.parent.mkdir(parents=True, exist_ok=True)
doc.save(stage, garbage=4, deflate=True, clean=True)
doc.close()
reader = PdfReader(str(stage))
writer = PdfWriter()
delete_indices = {op.page for op in delete_ops if op.page is not None}
order = reorder_ops[-1].pages if reorder_ops else None
indices = [i for i in range(len(reader.pages)) if i not in delete_indices]
if order:
indices = [i for i in order if i not in delete_indices]
for idx in indices:
writer.add_page(reader.pages[idx])
output_pdf.parent.mkdir(parents=True, exist_ok=True)
with output_pdf.open("wb") as f:
writer.write(f)
stage.unlink(missing_ok=True)
return output_pdf
|