Zhengshou Lai
5d0ed21e07
Add `mytoolkit pdf extract` to extract text from PDFs. Tries pdftotext (poppler) first for layout quality, falls back to pypdf.
262 lines
8.1 KiB
Python
262 lines
8.1 KiB
Python
"""PDF utilities."""
|
|
|
|
import shutil
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
import click
|
|
from pypdf import PdfReader, PdfWriter
|
|
from PIL import Image
|
|
|
|
from bin.md_to_pdf import expand_bookmarks
|
|
from bin.utils import handle_errors, run_command
|
|
|
|
|
|
@click.group()
|
|
def pdf():
|
|
"""PDF manipulation commands."""
|
|
pass
|
|
|
|
|
|
@pdf.command("compress")
|
|
@click.argument("files", nargs=-1, required=True)
|
|
@click.option(
|
|
"--quality",
|
|
"-q",
|
|
type=click.Choice(["screen", "ebook", "printer", "prepress"]),
|
|
default="screen",
|
|
help="Compression quality level",
|
|
)
|
|
@click.option(
|
|
"--bookmark-depth",
|
|
type=int,
|
|
default=1,
|
|
help="Expand PDF bookmarks to this depth after compression (1=no expansion).",
|
|
)
|
|
@click.option("--dry-run", "-n", is_flag=True, help="Show what would be done")
|
|
@handle_errors
|
|
def compress(files, quality, bookmark_depth, dry_run):
|
|
"""Compress PDF files using ghostscript."""
|
|
for pattern in files:
|
|
for pdf_file in Path(".").glob(pattern):
|
|
if pdf_file.suffix.lower() != ".pdf":
|
|
continue
|
|
|
|
output = pdf_file.with_suffix(".compressed.pdf")
|
|
cmd = [
|
|
"gs",
|
|
"-sDEVICE=pdfwrite",
|
|
"-dNOPAUSE",
|
|
"-dQUIET",
|
|
"-dBATCH",
|
|
f"-dPDFSETTINGS=/{quality}",
|
|
"-dCompatibilityLevel=1.4",
|
|
f"-sOutputFile={output}",
|
|
str(pdf_file),
|
|
]
|
|
|
|
if dry_run:
|
|
click.echo(f"Would compress: {pdf_file} -> {output}")
|
|
continue
|
|
|
|
click.echo(f"Compressing: {pdf_file}")
|
|
run_command(cmd, check=True)
|
|
output.replace(pdf_file)
|
|
if bookmark_depth > 1:
|
|
expand_bookmarks(pdf_file, bookmark_depth)
|
|
click.echo(f" Done: {pdf_file}")
|
|
|
|
|
|
@pdf.command("crop")
|
|
@click.argument("files", nargs=-1, required=True)
|
|
@click.option("--dry-run", "-n", is_flag=True, help="Show what would be done")
|
|
@handle_errors
|
|
def crop(files, dry_run):
|
|
"""Crop PDF files."""
|
|
for pattern in files:
|
|
for pdf_file in Path(".").glob(pattern):
|
|
if pdf_file.suffix.lower() != ".pdf":
|
|
continue
|
|
|
|
if dry_run:
|
|
click.echo(f"Would crop: {pdf_file}")
|
|
continue
|
|
|
|
click.echo(f"Cropping: {pdf_file}")
|
|
run_command(["pdfcrop", str(pdf_file), str(pdf_file)], check=False)
|
|
|
|
|
|
@pdf.command("to-tiff")
|
|
@click.argument("files", nargs=-1, required=True)
|
|
@click.option("--density", "-d", default="300", help="DPI density")
|
|
@click.option("--dry-run", "-n", is_flag=True, help="Show what would be done")
|
|
@handle_errors
|
|
def to_tiff(files, density, dry_run):
|
|
"""Convert PDF to TIFF."""
|
|
for pattern in files:
|
|
for pdf_file in Path(".").glob(pattern):
|
|
if pdf_file.suffix.lower() != ".pdf":
|
|
continue
|
|
output = pdf_file.with_suffix(".tiff")
|
|
|
|
if dry_run:
|
|
click.echo(f"Would convert: {pdf_file} -> {output}")
|
|
continue
|
|
|
|
click.echo(f"Converting: {pdf_file} -> {output}")
|
|
run_command(
|
|
["convert", "-density", density, str(pdf_file), "-compress", "zip", str(output)],
|
|
check=False,
|
|
)
|
|
|
|
|
|
@pdf.command("merge")
|
|
@click.argument("files", nargs=-1, required=True)
|
|
@click.option("--output", "-o", required=True, help="Output PDF file")
|
|
@click.option(
|
|
"--bookmark-depth",
|
|
type=int,
|
|
default=1,
|
|
help="Expand PDF bookmarks to this depth after merge (1=no expansion).",
|
|
)
|
|
@click.option("--dry-run", "-n", is_flag=True, help="Show what would be done")
|
|
@handle_errors
|
|
def merge_pdfs(files, output, bookmark_depth, dry_run):
|
|
"""Merge PDFs and/or images into a single PDF."""
|
|
writer = PdfWriter()
|
|
processed = []
|
|
|
|
for pattern in files:
|
|
for path in sorted(Path(".").glob(pattern)):
|
|
if path.suffix.lower() == ".pdf":
|
|
if dry_run:
|
|
click.echo(f"Would add PDF: {path}")
|
|
continue
|
|
reader = PdfReader(str(path))
|
|
for page in reader.pages:
|
|
writer.add_page(page)
|
|
processed.append(str(path))
|
|
elif path.suffix.lower() in (".png", ".jpg", ".jpeg"):
|
|
if dry_run:
|
|
click.echo(f"Would add image: {path}")
|
|
continue
|
|
img = Image.open(path)
|
|
if img.mode in ("RGBA", "P"):
|
|
img = img.convert("RGB")
|
|
# Save to temp PDF
|
|
temp_pdf = path.with_suffix(".temp.pdf")
|
|
img.save(temp_pdf, "PDF", resolution=150.0)
|
|
reader = PdfReader(str(temp_pdf))
|
|
writer.add_page(reader.pages[0])
|
|
temp_pdf.unlink()
|
|
processed.append(str(path))
|
|
|
|
if dry_run:
|
|
click.echo(f"Would write merged PDF to: {output}")
|
|
return
|
|
|
|
if not processed:
|
|
click.echo("No files found to merge")
|
|
return
|
|
|
|
with open(output, "wb") as f:
|
|
writer.write(f)
|
|
if bookmark_depth > 1:
|
|
expand_bookmarks(Path(output), bookmark_depth)
|
|
click.echo(f"Merged {len(processed)} files into {output}")
|
|
|
|
|
|
@pdf.command("extract")
|
|
@click.argument("input_file", type=click.Path(exists=True, dir_okay=False, path_type=Path))
|
|
@click.option(
|
|
"-o",
|
|
"--output",
|
|
type=click.Path(dir_okay=False, path_type=Path),
|
|
help="Output text file. Defaults to stdout.",
|
|
)
|
|
@click.option(
|
|
"-p",
|
|
"--pages",
|
|
help="Page range to extract (e.g. 1-5, 1,3,5). Defaults to all pages.",
|
|
)
|
|
@click.option(
|
|
"--layout/--no-layout",
|
|
default=False,
|
|
help="Preserve original layout using pdftotext -layout. Only applies when pdftotext is available.",
|
|
)
|
|
@handle_errors
|
|
def extract_text(input_file, output, pages, layout):
|
|
"""Extract text from a PDF file.
|
|
|
|
Tries pdftotext (poppler) first for best quality, falls back to pypdf.
|
|
"""
|
|
text = ""
|
|
use_pdftotext = shutil.which("pdftotext") is not None
|
|
|
|
if use_pdftotext:
|
|
cmd = ["pdftotext"]
|
|
if layout:
|
|
cmd.append("-layout")
|
|
if pages:
|
|
cmd.extend(["-f", str(pages.split("-")[0]), "-l", str(pages.split("-")[-1])])
|
|
cmd.extend([str(input_file), "-"])
|
|
result = subprocess.run(cmd, capture_output=True, text=True)
|
|
text = result.stdout
|
|
else:
|
|
reader = PdfReader(str(input_file))
|
|
page_indices = _parse_page_range(pages, len(reader.pages))
|
|
text = "\n\n".join(
|
|
reader.pages[i].extract_text() or ""
|
|
for i in page_indices
|
|
)
|
|
|
|
if output:
|
|
output.write_text(text, encoding="utf-8")
|
|
click.echo(f"Extracted text to {output}")
|
|
else:
|
|
click.echo(text)
|
|
|
|
|
|
def _parse_page_range(pages: str | None, total: int) -> list[int]:
|
|
"""Parse a page range string into 0-based page indices."""
|
|
if not pages:
|
|
return list(range(total))
|
|
|
|
indices = set()
|
|
for part in pages.split(","):
|
|
part = part.strip()
|
|
if "-" in part:
|
|
start, end = part.split("-", 1)
|
|
start = int(start.strip()) - 1 if start.strip() else 0
|
|
end = int(end.strip()) if end.strip() else total
|
|
indices.update(range(max(0, start), min(total, end)))
|
|
else:
|
|
idx = int(part) - 1
|
|
if 0 <= idx < total:
|
|
indices.add(idx)
|
|
return sorted(indices)
|
|
|
|
|
|
@pdf.command("bookmark")
|
|
@click.argument("input_file", type=click.Path(exists=True, dir_okay=False, path_type=Path))
|
|
@click.option(
|
|
"--depth",
|
|
type=int,
|
|
required=True,
|
|
help="Bookmark expansion depth. 1=top-level only, 2=expand to second level, etc.",
|
|
)
|
|
@click.option(
|
|
"-o",
|
|
"--output",
|
|
type=click.Path(dir_okay=False, path_type=Path),
|
|
help="Output PDF path. Defaults to overwriting the input file.",
|
|
)
|
|
@handle_errors
|
|
def bookmark(input_file, depth, output):
|
|
"""Set PDF bookmark outline expansion depth."""
|
|
target = output if output else input_file
|
|
if output and output.resolve() != input_file.resolve():
|
|
shutil.copy2(str(input_file), str(target))
|
|
expand_bookmarks(target, depth)
|
|
click.echo(f"Bookmarks expanded to depth {depth}: {target}")
|