pryscraper/parser.py

"""Pry — document parser for PDF, DOCX, images, CSV, JSON, HTML.
Uses asyncio subprocess for CPU-bound OCR to avoid blocking the event loop.
Temp files are always cleaned up in finally blocks.
"""

import asyncio
import io
import os
import tempfile
from pathlib import Path
from typing import Any


class DocumentParser:
    """Parse PDFs, DOCX, images, and other document formats to clean text.
    CPU-bound operations (OCR, PDF parsing) run in executor threads."""

    async def parse(self, url: str, timeout: int = 60) -> dict[str, Any]:
        import httpx

        async with httpx.AsyncClient(
            timeout=httpx.Timeout(timeout), follow_redirects=True
        ) as client:
            resp = await client.get(url)
            resp.raise_for_status()
            content = resp.content
            content_type = resp.headers.get("content-type", "")
            filename = Path(url).name.lower()

        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self._parse_bytes, content, content_type, filename)

    def _parse_bytes(self, data: bytes, content_type: str, filename: str = "") -> dict[str, Any]:
        if filename.endswith(".pdf") or "pdf" in content_type:
            return self._parse_pdf(data)
        elif filename.endswith(".docx") or "word" in content_type:
            return self._parse_docx(data)
        elif filename.endswith((".md", ".markdown")) or "markdown" in content_type:
            return {
                "text": data.decode("utf-8", errors="replace"),
                "format": "markdown",
                "pages": 1,
            }
        elif filename.endswith(".csv") or "csv" in content_type:
            return {"text": data.decode("utf-8", errors="replace"), "format": "csv", "pages": 1}
        elif filename.endswith(".json") or "json" in content_type:
            return {"text": data.decode("utf-8", errors="replace"), "format": "json", "pages": 1}
        elif filename.endswith((".html", ".htm")) or "html" in content_type:
            import trafilatura

            text = data.decode("utf-8", errors="replace")
            extracted = trafilatura.extract(text, output_format="markdown")
            return {"text": extracted or text[:500000], "format": "html", "pages": 1}
        elif (
            filename.endswith((".png", ".jpg", ".jpeg", ".gif", ".webp")) or "image" in content_type
        ):
            return self._parse_image(data)
        elif filename.endswith(".txt") or "text" in content_type:
            return {"text": data.decode("utf-8", errors="replace"), "format": "text", "pages": 1}
        else:
            try:
                text = data.decode("utf-8", errors="replace")
                return {"text": text[:500000], "format": "unknown", "pages": 1}
            except UnicodeDecodeError:
                return {
                    "text": f"[Binary file: {filename}, {len(data)} bytes]",
                    "format": "binary",
                    "pages": 0,
                }

    def _parse_pdf(self, data: bytes) -> dict[str, Any]:
        from pypdf import PdfReader

        reader = PdfReader(io.BytesIO(data))
        pages = []
        for page in reader.pages:
            text = page.extract_text()
            if text:
                pages.append(text)
        return {
            "text": "\n\n".join(pages),
            "format": "pdf",
            "pages": len(pages),
            "metadata": {
                "title": reader.metadata.get("/Title", "") or "",
                "author": reader.metadata.get("/Author", "") or "",
            },
        }

    def _parse_docx(self, data: bytes) -> dict[str, Any]:
        from docx import Document

        doc = Document(io.BytesIO(data))
        paragraphs = [p.text for p in doc.paragraphs if p.text.strip()]
        return {"text": "\n\n".join(paragraphs), "format": "docx", "pages": 1}

    def _parse_image(self, data: bytes) -> dict[str, Any]:
        """OCR via tesseract. Temp file cleaned up in finally block."""
        fname = None
        try:
            from PIL import Image

            img = Image.open(io.BytesIO(data))
            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
                img.save(f, format="PNG")
                fname = f.name
            result = __import__("subprocess").run(
                ["tesseract", fname, "stdout", "--oem", "1", "-l", "eng"],
                capture_output=True,
                text=True,
                timeout=30,
            )
            return {"text": result.stdout, "format": "image", "pages": 1}
        except Exception as e:
            return {"text": f"[Image OCR failed: {e}]", "format": "image", "pages": 0}
        finally:
            if fname and os.path.exists(fname):
                try:
                    os.unlink(fname)
                except Exception:
                    pass