compress_pdfs.py

"""
Utility script to compress PDF files recursively using PyPDF.

The script walks every PDF in the provided folder (you can also pick a folder
via a dialog), reduces the quality of embedded images using PyPDF's
``image.replace`` helper, applies lossless Flate compression to each page's
content stream, and overwrites the original PDF only when the compressed
version is smaller. A progress report and aggregate savings are printed to
stdout.
"""

from __future__ import annotations

import argparse
import logging
import sys
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

from PIL import UnidentifiedImageError
from pypdf import PdfWriter
from pypdf.errors import PdfReadError

try:
    import tkinter as tk
    from tkinter import filedialog
except Exception:  # pragma: no cover - tkinter may be unavailable on some systems
    tk = None
    filedialog = None

for logger_name in ("pypdf", "pypdf._page", "pypdf._xobj_image_helpers"):
    logging.getLogger(logger_name).setLevel(logging.ERROR)


@dataclass
class FileCompressionStats:
    """Stores size statistics for a single PDF file."""

    before_bytes: int
    after_bytes: int
    images_processed: int
    images_skipped: int
    pages_compressed: int
    processed: bool
    message: str = ""

    @property
    def saved_bytes(self) -> int:
        return self.before_bytes - self.after_bytes


def parse_args(argv: Iterable[str]) -> argparse.Namespace:
    argv = list(argv)
    legacy_folder: Path | None = None
    if argv and not argv[0].startswith("-"):
        potential = argv.pop(0)
        try:
            float(potential)
        except ValueError:
            legacy_folder = Path(potential)
        else:
            argv.insert(0, potential)

    parser = argparse.ArgumentParser(
        description=(
            "Compress PDF files under a folder by reducing image quality and "
            "applying lossless Flate compression using PyPDF."
        )
    )
    parser.add_argument(
        "--folder",
        "-f",
        type=Path,
        default=legacy_folder,
        help=(
            "Root folder to search for PDF files recursively. "
            "When omitted, a folder picker dialog is opened."
        ),
    )
    parser.add_argument(
        "compression_ratio",
        type=float,
        help=(
            "Compression ratio between 0 and 1. Values closer to 0 apply stronger "
            "compression by lowering image quality more aggressively."
        ),
    )
    parser.add_argument(
        "--lossless-level",
        type=int,
        default=None,
        help=(
            "Optional override for the lossless (Flate) compression level "
            "applied to each page. Must be between 1 (fastest) and 9 (highest). "
            "When omitted, the level is derived from the compression ratio."
        ),
    )
    parsed = parser.parse_args(argv)

    if parsed.compression_ratio <= 0:
        parser.error("compression_ratio must be greater than 0.")
    if parsed.compression_ratio > 1:
        print(
            "Compression ratio is greater than 1; capping to 1 (no extra quality loss).",
            file=sys.stderr,
        )
        parsed.compression_ratio = 1.0

    if parsed.lossless_level is not None and not (1 <= parsed.lossless_level <= 9):
        parser.error("--lossless-level must be between 1 and 9.")

    if parsed.folder is None:
        parsed.folder = prompt_for_folder_via_dialog(parser)

    parsed.folder = parsed.folder.expanduser().resolve()
    return parsed


def prompt_for_folder_via_dialog(parser: argparse.ArgumentParser) -> Path:
    """
    Ask the user to select a folder via a file explorer dialog.

    Falls back to an error if tkinter is unavailable or the user cancels.
    """
    if tk is None or filedialog is None:
        parser.error(
            "tkinter is not available, please provide the folder path explicitly."
        )

    root = tk.Tk()
    root.withdraw()
    folder = filedialog.askdirectory(title="Select folder containing PDFs")
    root.destroy()

    if not folder:
        parser.error("No folder selected; please provide the folder path.")

    return Path(folder)


def format_bytes(num_bytes: int) -> str:
    """Human readable byte formatter."""
    if num_bytes < 1024:
        return f"{num_bytes} B"
    units = ["KB", "MB", "GB", "TB"]
    size = float(num_bytes)
    for unit in units:
        size /= 1024.0
        if size < 1024.0:
            return f"{size:.2f} {unit}"
    return f"{size:.2f} PB"


def quality_from_ratio(ratio: float) -> int:
    """
    Map the ratio (0-1) to a JPEG quality value between 10 and 95.

    Smaller ratios translate to lower qualities, capping the range to avoid
    unreadable images or redundant recompression when quality is already high.
    """
    min_quality, max_quality = 10, 95
    computed = min_quality + int((max_quality - min_quality) * ratio)
    return max(min_quality, min(max_quality, computed))


def lossless_level_from_ratio(ratio: float) -> int:
    """
    Map the ratio (0-1) to a Flate compression level between 1 and 9.

    Lower ratios produce a higher compression level (slower but smaller).
    """
    return max(1, min(9, 1 + int((1 - ratio) * 8)))


def compress_pdf(
    pdf_path: Path,
    ratio: float,
    lossless_level_override: int | None,
) -> FileCompressionStats:
    """Compress a single PDF file using PyPDF and return its statistics."""
    before_size = pdf_path.stat().st_size

    writer = PdfWriter(clone_from=str(pdf_path))

    target_quality = quality_from_ratio(ratio)
    lossless_level = lossless_level_override or lossless_level_from_ratio(ratio)

    images_processed = 0
    images_skipped = 0
    pages_compressed = 0

    for page_index, page in enumerate(writer.pages):
        try:
            image_ids = page._get_ids_image()
        except Exception as exc:
            print(f"      ! Failed to enumerate image references on page {page_index + 1}: {exc}")
            image_ids = []

        for image_id in image_ids:
            try:
                image_obj = page._get_image(image_id)
            except (UnidentifiedImageError, PdfReadError, KeyError, ValueError) as exc:
                images_skipped += 1
                continue
            except Exception as exc:
                print(
                    f"      ! Failed to access image {image_id!r} on page {page_index + 1}: {exc}"
                )
                images_skipped += 1
                continue

            pil_image = image_obj.image
            if pil_image is None:
                images_skipped += 1
                continue

            try:
                image_obj.replace(pil_image, quality=target_quality)
                images_processed += 1
            except Exception as exc:
                print(
                    f"      ! Failed to adjust image {image_obj.name!r} on page {page_index + 1}: {exc}"
                )
                images_skipped += 1

        try:
            page.compress_content_streams(level=lossless_level)
            pages_compressed += 1
        except Exception as exc:
            print(
                f"      ! Failed to apply lossless compression on page {page_index + 1}: {exc}"
            )

    temp_file = tempfile.NamedTemporaryFile(
        suffix=".pdf", delete=False, dir=str(pdf_path.parent)
    )
    temp_path = Path(temp_file.name)
    with temp_file:
        try:
            writer.write(temp_file)
        except Exception as exc:
            writer.close()
            temp_path.unlink(missing_ok=True)
            return FileCompressionStats(
                before_size,
                before_size,
                images_processed,
                images_skipped,
                pages_compressed,
                False,
                f"Failed to write compressed file: {exc}",
            )

    writer.close()

    after_size = temp_path.stat().st_size if temp_path.exists() else before_size

    if after_size >= before_size or after_size == 0:
        temp_path.unlink(missing_ok=True)
        return FileCompressionStats(
            before_size,
            before_size,
            images_processed,
            images_skipped,
            pages_compressed,
            False,
            "Compressed file was not smaller; original kept.",
        )

    try:
        temp_path.replace(pdf_path)
    except Exception as exc:
        temp_path.unlink(missing_ok=True)
        return FileCompressionStats(
            before_size,
            before_size,
            images_processed,
            images_skipped,
            pages_compressed,
            False,
            f"Failed to overwrite original PDF: {exc}",
        )

    return FileCompressionStats(
        before_size,
        after_size,
        images_processed,
        images_skipped,
        pages_compressed,
        True,
        "Compressed successfully.",
    )


def main(argv: Iterable[str]) -> int:
    args = parse_args(argv)

    if not args.folder.exists():
        print(f"Folder does not exist: {args.folder}", file=sys.stderr)
        return 1

    pdf_files = sorted(args.folder.rglob("*.pdf"))
    if not pdf_files:
        print(f"No PDF files found under {args.folder}")
        return 0

    print(f"Scanning: {args.folder}")
    print(f"Found {len(pdf_files)} PDF file(s). Starting compression...")

    total_before = 0
    total_after = 0
    total_images = 0
    total_pages = 0
    total_skipped_images = 0
    files_compressed = 0
    files_skipped = 0

    for index, pdf_file in enumerate(pdf_files, start=1):
        rel_path = pdf_file.relative_to(args.folder)
        print(f"[{index}/{len(pdf_files)}] Processing {rel_path} ...")
        stats = compress_pdf(pdf_file, args.compression_ratio, args.lossless_level)
        total_before += stats.before_bytes
        total_after += stats.after_bytes
        total_images += stats.images_processed
        total_pages += stats.pages_compressed
        total_skipped_images += stats.images_skipped

        if not stats.processed:
            files_skipped += 1
            skip_note = stats.message
            if stats.images_skipped:
                skip_note += f" (images skipped: {stats.images_skipped})"
            print(f"    Skipped: {skip_note}")
            continue

        files_compressed += 1
        change = stats.saved_bytes
        if change > 0:
            print(
                f"    {stats.message} Images adjusted: {stats.images_processed}, "
                f"images skipped: {stats.images_skipped}, "
                f"pages recompressed: {stats.pages_compressed}. "
                f"Saved {format_bytes(change)}."
            )
        elif change == 0:
            print(
                f"    {stats.message} Images adjusted: {stats.images_processed}, "
                f"images skipped: {stats.images_skipped}, "
                f"pages recompressed: {stats.pages_compressed}. Size unchanged."
            )
        else:
            print(
                f"    {stats.message} Images adjusted: {stats.images_processed}, "
                f"images skipped: {stats.images_skipped}, "
                f"pages recompressed: {stats.pages_compressed}. "
                f"Grew by {format_bytes(abs(change))}."
            )

    print("\nCompression complete.")
    print(f"Total PDFs processed: {len(pdf_files)}")
    print(f"Files compressed    : {files_compressed}")
    print(f"Files skipped       : {files_skipped}")
    print(f"Images adjusted     : {total_images}")
    print(f"Images skipped      : {total_skipped_images}")
    print(f"Pages recompressed  : {total_pages}")
    print(f"Initial total size  : {format_bytes(total_before)}")
    print(f"Final total size    : {format_bytes(total_after)}")

    savings = total_before - total_after
    if savings > 0:
        print(f"Overall savings     : {format_bytes(savings)}")
    elif savings == 0:
        print("Overall savings     : none")
    else:
        print(f"Overall size increase: {format_bytes(abs(savings))}")

    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))