gauravmeena0708/epfo-circulars

1"""Core in-memory PDF manipulation utilities for merging, splitting, rotating,2compressing, format conversion, and OCR processing.3"""4 5from __future__ import annotations6 7import io8import os9import re10from typing import Sequence11import numpy as np12from PIL import Image13 14try:15    import pymupdf as fitz16except ImportError:  # pragma: no cover17    import fitz18 19 20class PDFProcessingError(ValueError):21    """Raised when PDF processing fails due to invalid input, encryption, or corruption."""22 23 24def _open_pdf(pdf_data: bytes) -> fitz.Document:25    """Safely opens an in-memory PDF, checking for corruption and encryption."""26    if not pdf_data:27        raise PDFProcessingError("The provided PDF data is empty.")28    try:29        doc = fitz.open(stream=pdf_data, filetype="pdf")30    except Exception as exc:31        raise PDFProcessingError(f"Could not open or parse the PDF document: {exc}") from exc32 33    if doc.needs_pass:34        doc.close()35        raise PDFProcessingError("The PDF document is password-protected. Please upload an unlocked PDF.")36    return doc37 38 39def parse_page_ranges(range_str: str, max_pages: int) -> list[int]:40    """41    Parses a page range string (1-indexed) into a list of 0-indexed page indices.42    Supports formats: "1", "1-3", "1, 3, 5", "2-4, 1", "all".43    """44    if not range_str or not range_str.strip():45        raise PDFProcessingError("Page range cannot be empty.")46 47    clean_str = range_str.strip().lower()48    if clean_str == "all":49        return list(range(max_pages))50 51    indices: list[int] = []52    tokens = [t.strip() for t in clean_str.split(",") if t.strip()]53    if not tokens:54        raise PDFProcessingError("Invalid page range format.")55 56    for token in tokens:57        if "-" in token:58            parts = token.split("-")59            if len(parts) != 2:60                raise PDFProcessingError(f"Invalid range format: '{token}'")61            try:62                start = int(parts[0].strip())63                end = int(parts[1].strip())64            except ValueError:65                raise PDFProcessingError(f"Invalid numbers in range: '{token}'")66 67            if start < 1 or end > max_pages or start > end:68                raise PDFProcessingError(69                    f"Range '{start}-{end}' is invalid. Document has {max_pages} pages (allowed: 1-{max_pages})."70                )71            indices.extend(range(start - 1, end))72        else:73            try:74                page_num = int(token)75            except ValueError:76                raise PDFProcessingError(f"Invalid page number: '{token}'")77 78            if page_num < 1 or page_num > max_pages:79                raise PDFProcessingError(80                    f"Page number {page_num} is out of bounds (1-{max_pages})."81                )82            indices.append(page_num - 1)83 84    return indices85 86 87def merge_pdfs(pdf_bytes_list: Sequence[bytes]) -> bytes:88    """89    Merges multiple PDF files in sequential order.90    Returns the merged PDF as bytes.91    """92    if not pdf_bytes_list:93        raise PDFProcessingError("At least one PDF file is required to merge.")94 95    merged_doc = fitz.open()96    try:97        for index, pdf_data in enumerate(pdf_bytes_list):98            sub_doc = _open_pdf(pdf_data)99            try:100                merged_doc.insert_pdf(sub_doc)101            finally:102                sub_doc.close()103 104        if len(merged_doc) == 0:105            raise PDFProcessingError("The resulting merged PDF has 0 pages.")106 107        return merged_doc.tobytes(deflate=True, garbage=3)108    finally:109        merged_doc.close()110 111 112def split_pdf(pdf_bytes: bytes, page_range_str: str) -> bytes:113    """114    Extracts specified pages from a PDF and returns a new PDF containing only those pages.115    """116    src_doc = _open_pdf(pdf_bytes)117    try:118        target_indices = parse_page_ranges(page_range_str, len(src_doc))119        out_doc = fitz.open()120        try:121            for page_idx in target_indices:122                out_doc.insert_pdf(src_doc, from_page=page_idx, to_page=page_idx)123            return out_doc.tobytes(deflate=True, garbage=3)124        finally:125            out_doc.close()126    finally:127        src_doc.close()128 129 130def split_pdf_to_individual_pages(pdf_bytes: bytes) -> list[tuple[str, bytes]]:131    """132    Splits a multi-page PDF into a list of single-page PDFs: [('page_1.pdf', bytes), ...].133    """134    src_doc = _open_pdf(pdf_bytes)135    try:136        pages_list: list[tuple[str, bytes]] = []137        for i in range(len(src_doc)):138            single_doc = fitz.open()139            try:140                single_doc.insert_pdf(src_doc, from_page=i, to_page=i)141                page_bytes = single_doc.tobytes(deflate=True, garbage=3)142                pages_list.append((f"page_{i + 1}.pdf", page_bytes))143            finally:144                single_doc.close()145        return pages_list146    finally:147        src_doc.close()148 149 150def rotate_pdf_pages(151    pdf_bytes: bytes,152    rotation_degrees: int,153    page_range_str: str | None = None,154) -> bytes:155    """156    Rotates pages by the specified degrees (90, 180, 270).157    If page_range_str is provided, only specified pages are rotated.158    """159    if rotation_degrees not in (90, 180, 270):160        raise PDFProcessingError(161            f"Invalid rotation degrees: {rotation_degrees}. Allowed values: 90, 180, 270."162        )163 164    doc = _open_pdf(pdf_bytes)165    try:166        if page_range_str and page_range_str.strip():167            target_indices = set(parse_page_ranges(page_range_str, len(doc)))168        else:169            target_indices = set(range(len(doc)))170 171        for idx in target_indices:172            page = doc[idx]173            new_rotation = (page.rotation + rotation_degrees) % 360174            page.set_rotation(new_rotation)175 176        return doc.tobytes(deflate=True, garbage=3)177    finally:178        doc.close()179 180 181def compress_pdf(182    pdf_bytes: bytes,183    deflate: bool = True,184    clean: bool = True,185) -> tuple[bytes, int, int]:186    """187    Compresses and optimizes a PDF by deflating streams and cleaning redundant objects.188    Returns: (compressed_bytes, original_size_bytes, new_size_bytes)189    """190    doc = _open_pdf(pdf_bytes)191    try:192        orig_size = len(pdf_bytes)193        compressed_bytes = doc.tobytes(194            garbage=4,195            deflate=deflate,196            clean=clean,197            linear=True,198        )199        return compressed_bytes, orig_size, len(compressed_bytes)200    finally:201        doc.close()202 203 204def images_to_pdf(images_data: Sequence[tuple[str, bytes]]) -> bytes:205    """206    Combines multiple images into a single PDF document.207    images_data: list of (filename, image_bytes)208    """209    if not images_data:210        raise PDFProcessingError("At least one image is required to build a PDF.")211 212    pil_images: list[Image.Image] = []213    try:214        for name, img_bytes in images_data:215            if not img_bytes:216                continue217            try:218                img = Image.open(io.BytesIO(img_bytes))219                # Convert RGBA/P/LA to RGB for clean PDF output220                if img.mode in ("RGBA", "P", "LA"):221                    rgb_img = Image.new("RGB", img.size, (255, 255, 255))222                    if img.mode == "RGBA":223                        rgb_img.paste(img, mask=img.split()[3])224                    else:225                        rgb_img.paste(img.convert("RGBA"), mask=img.convert("RGBA").split()[3])226                    pil_images.append(rgb_img)227                elif img.mode != "RGB":228                    pil_images.append(img.convert("RGB"))229                else:230                    pil_images.append(img)231            except Exception as exc:232                raise PDFProcessingError(f"Could not load image '{name}': {exc}") from exc233 234        if not pil_images:235            raise PDFProcessingError("No valid images could be decoded.")236 237        out_buf = io.BytesIO()238        pil_images[0].save(239            out_buf,240            format="PDF",241            save_all=True,242            append_images=pil_images[1:] if len(pil_images) > 1 else [],243            resolution=150.0,244        )245        return out_buf.getvalue()246    finally:247        for img in pil_images:248            img.close()249 250 251def pdf_to_images(252    pdf_bytes: bytes,253    dpi: int = 200,254    image_format: str = "PNG",255) -> list[tuple[str, bytes]]:256    """257    Renders each page of a PDF as a high-resolution image.258    Returns: [('page_1.png', image_bytes), ...]259    """260    doc = _open_pdf(pdf_bytes)261    ext = "jpg" if image_format.upper() in ("JPG", "JPEG") else "png"262    fmt = "jpeg" if ext == "jpg" else "png"263 264    try:265        rendered: list[tuple[str, bytes]] = []266        scale = max(dpi, 72) / 72267        matrix = fitz.Matrix(scale, scale)268 269        for i, page in enumerate(doc):270            pix = page.get_pixmap(matrix=matrix, alpha=False)271            img_bytes = pix.tobytes(output=fmt)272            rendered.append((f"page_{i + 1}.{ext}", img_bytes))273 274        return rendered275    finally:276        doc.close()277 278 279def ocr_scanned_pdf(280    pdf_bytes: bytes,281    ocr_reader,282    dpi: int = 200,283) -> tuple[str, bytes]:284    """285    Runs OCR on each page of a scanned PDF.286    Returns: (extracted_text, searchable_pdf_bytes)287    """288    doc = _open_pdf(pdf_bytes)289    extracted_pages: list[str] = []290    scale = max(dpi, 72) / 72291    matrix = fitz.Matrix(scale, scale)292 293    try:294        out_doc = fitz.open()295        for i, page in enumerate(doc):296            pix = page.get_pixmap(matrix=matrix, alpha=False)297            image_np = np.frombuffer(pix.samples, dtype=np.uint8).reshape(298                pix.height,299                pix.width,300                pix.n,301            )302 303            # Read text with EasyOCR304            results = ocr_reader.readtext(image_np, detail=0, paragraph=True)305            page_text = "\n".join(str(r).strip() for r in results if str(r).strip())306            extracted_pages.append(f"--- [Page {i + 1} / {len(doc)}] ---\n{page_text}")307 308            # Recreate page in output PDF with OCR text309            new_page = out_doc.new_page(width=page.rect.width, height=page.rect.height)310            new_page.insert_image(page.rect, stream=pix.tobytes(output="png"))311            if page_text:312                # Insert hidden/overlay text to make it searchable313                new_page.insert_text(314                    (36, 36),315                    page_text,316                    fontsize=8,317                    color=(0, 0, 0),318                    render_mode=3,  # Invisible text render mode in PDF319                )320 321        full_extracted_text = "\n\n".join(extracted_pages)322        searchable_pdf_bytes = out_doc.tobytes(deflate=True, garbage=3)323        out_doc.close()324        return full_extracted_text, searchable_pdf_bytes325    finally:326        doc.close()327