paperless-ngx/src/documents/parsers.py

from __future__ import annotations

import logging
import math
import mimetypes
import shutil
import subprocess
import tempfile
from pathlib import Path
from typing import TYPE_CHECKING

from django.conf import settings

from documents.loggers import LoggingMixin
from documents.utils import copy_file_with_basic_stats
from documents.utils import run_subprocess
from paperless.parsers.registry import get_parser_registry

if TYPE_CHECKING:
    import datetime

logger = logging.getLogger("paperless.parsing")


def is_mime_type_supported(mime_type: str) -> bool:
    """
    Returns True if the mime type is supported, False otherwise
    """
    return get_parser_registry().get_parser_for_file(mime_type, "") is not None


def get_default_file_extension(mime_type: str) -> str:
    """
    Returns the default file extension for a mimetype, or
    an empty string if it could not be determined
    """
    parser_class = get_parser_registry().get_parser_for_file(mime_type, "")
    if parser_class is not None:
        supported = parser_class.supported_mime_types()
        if mime_type in supported:
            return supported[mime_type]

    ext = mimetypes.guess_extension(mime_type)
    return ext if ext else ""


def is_file_ext_supported(ext: str) -> bool:
    """
    Returns True if the file extension is supported, False otherwise
    TODO: Investigate why this really exists, why not use mimetype
    """
    if ext:
        return ext.lower() in get_supported_file_extensions()
    else:
        return False


def get_supported_file_extensions() -> set[str]:
    extensions = set()
    for parser_class in get_parser_registry().all_parsers():
        for mime_type, ext in parser_class.supported_mime_types().items():
            extensions.update(mimetypes.guess_all_extensions(mime_type))
            # Python's stdlib might be behind, so also add what the parser
            # says is the default extension
            # This makes image/webp supported on Python < 3.11
            extensions.add(ext)

    return extensions


def get_default_thumbnail() -> Path:
    """
    Returns the path to a generic thumbnail
    """
    return (Path(__file__).parent / "resources" / "document.webp").resolve()


_THUMBNAIL_MAX_WIDTH = 500
_THUMBNAIL_MAX_HEIGHT = 5000
# Used only when the page geometry cannot be read
_THUMBNAIL_FALLBACK_DPI = 150
# Applied before supersampling, so tiny pages are not enlarged
_THUMBNAIL_MAX_DPI = 300
# Rendering at a multiple and downsampling keeps text crisper
_THUMBNAIL_SUPERSAMPLE = 2


def rasterize_pdf_page_to_png(
    in_path: Path,
    out_path: Path,
    *,
    dpi: int,
    logging_group=None,
) -> None:
    """
    Rasterizes the first page of a PDF to a PNG with pdftoppm.
    """
    # -singlefile drops the page number and -png appends ".png", so pass the
    # path without its suffix
    args = [
        "pdftoppm",
        "-f",
        "1",
        "-l",
        "1",
        "-r",
        str(dpi),
        "-png",
        "-singlefile",
        "-cropbox",
        str(in_path),
        str(out_path.with_suffix("")),
    ]

    logger.debug("Execute: " + " ".join(args), extra={"group": logging_group})

    try:
        run_subprocess(args, logger=logger)
    except subprocess.CalledProcessError as e:
        raise ParseError(f"pdftoppm failed at {args}") from e
    except Exception as e:  # pragma: no cover
        raise ParseError("Unknown error running pdftoppm") from e


def encode_thumbnail_webp(
    png_path: Path,
    out_path: Path,
    *,
    supersample: int = 1,
) -> None:
    """
    Flattens alpha onto white, undoes supersampling, shrinks to fit and saves as WebP.
    """
    from PIL import Image

    try:
        with Image.open(png_path) as im:
            if im.mode in ("RGBA", "LA"):
                flattened = Image.new("RGB", im.size, (255, 255, 255))
                flattened.paste(im, mask=im.split()[-1])
            else:
                flattened = im.convert("RGB")

        if supersample > 1:
            flattened = flattened.resize(
                (
                    max(1, round(flattened.width / supersample)),
                    max(1, round(flattened.height / supersample)),
                ),
                Image.Resampling.LANCZOS,
            )

        flattened.thumbnail((_THUMBNAIL_MAX_WIDTH, _THUMBNAIL_MAX_HEIGHT))
        flattened.save(out_path, format="WEBP")
    except (OSError, Image.DecompressionBombError) as e:
        raise ParseError(f"Unable to encode thumbnail from {png_path}") from e


def _compute_thumbnail_dpi(in_path: Path, logging_group=None) -> tuple[int, int]:
    """
    Returns (dpi, supersample). Unknown geometry is not supersampled, since
    the render size cannot be bounded.
    """
    from paperless.parsers.utils import get_pdf_first_page_size_points

    size = get_pdf_first_page_size_points(in_path)
    if size is None:
        logger.debug(
            "Could not read PDF page size, using fallback DPI",
            extra={"group": logging_group},
        )
        return _THUMBNAIL_FALLBACK_DPI, 1

    width_pts, height_pts = size
    dpi_for_width = _THUMBNAIL_MAX_WIDTH * 72 / width_pts
    dpi_for_height = _THUMBNAIL_MAX_HEIGHT * 72 / height_pts
    # Round up so the downsampled render is never a few pixels short of the
    # target; the shrink-only clamp in encode_thumbnail_webp trims the excess.
    dpi = max(
        1,
        math.ceil(min(_THUMBNAIL_MAX_DPI, dpi_for_width, dpi_for_height)),
    )
    # At the 1 DPI floor the page is already oversized, so do not supersample
    return dpi, 1 if dpi == 1 else _THUMBNAIL_SUPERSAMPLE


def _render_pdf_thumbnail(
    in_path: Path,
    png_path: Path,
    out_path: Path,
    logging_group=None,
) -> None:
    dpi, supersample = _compute_thumbnail_dpi(in_path, logging_group=logging_group)
    rasterize_pdf_page_to_png(
        in_path,
        png_path,
        dpi=dpi * supersample,
        logging_group=logging_group,
    )
    encode_thumbnail_webp(png_path, out_path, supersample=supersample)


def _repair_pdf_with_qpdf(in_path: Path, out_path: Path) -> None:
    # qpdf exits 3 after a repair; --warning-exit-0 keeps that from failing
    try:
        shutil.copy(in_path, out_path)
        run_subprocess(
            ["qpdf", "--warning-exit-0", "--replace-input", str(out_path)],
            logger=logger,
        )
    except (subprocess.CalledProcessError, OSError) as e:
        raise ParseError(f"qpdf repair failed for {in_path}") from e


def make_thumbnail_from_pdf_qpdf_fallback(
    in_path: Path,
    temp_dir: Path,
    logging_group=None,
) -> Path:
    png_path = temp_dir / "page1_repaired.png"
    out_path = temp_dir / "convert_qpdf.webp"
    repaired_path = temp_dir / "repaired.pdf"

    logger.warning(
        "Thumbnail generation with pdftoppm failed, attempting qpdf repair and retry.",
        extra={"group": logging_group},
    )

    try:
        _repair_pdf_with_qpdf(in_path, repaired_path)
        _render_pdf_thumbnail(repaired_path, png_path, out_path, logging_group)
        return out_path

    except ParseError as e:
        logger.error(f"Unable to make thumbnail after qpdf repair: {e}")
        # The caller might expect a generated thumbnail that can be moved,
        # so we need to copy it before it gets moved.
        # https://github.com/paperless-ngx/paperless-ngx/issues/3631
        default_thumbnail_path = temp_dir / "document.webp"
        copy_file_with_basic_stats(get_default_thumbnail(), default_thumbnail_path)
        return default_thumbnail_path


def make_thumbnail_from_pdf(in_path: Path, temp_dir: Path, logging_group=None) -> Path:
    """
    The thumbnail of a PDF is just a 500px wide image of the first page.
    """
    png_path: Path = temp_dir / "page1.png"
    out_path: Path = temp_dir / "convert.webp"

    try:
        _render_pdf_thumbnail(in_path, png_path, out_path, logging_group)
    except ParseError as e:
        logger.error(f"Unable to make thumbnail with pdftoppm: {e}")
        out_path = make_thumbnail_from_pdf_qpdf_fallback(
            in_path,
            temp_dir,
            logging_group,
        )

    return out_path


class ParseError(Exception):
    pass


class DocumentParser(LoggingMixin):
    """
    Subclass this to make your own parser.  Have a look at
    `paperless_tesseract.parsers` for inspiration.
    """

    logging_name = "paperless.parsing"

    def __init__(self, logging_group, progress_callback=None) -> None:
        super().__init__()
        self.renew_logging_group()
        self.logging_group = logging_group
        self.settings = self.get_settings()
        settings.SCRATCH_DIR.mkdir(parents=True, exist_ok=True)
        self.tempdir = Path(
            tempfile.mkdtemp(prefix="paperless-", dir=settings.SCRATCH_DIR),
        )

        self.archive_path = None
        self.text = None
        self.date: datetime.datetime | None = None
        self.progress_callback = progress_callback

    def progress(self, current_progress, max_progress) -> None:
        if self.progress_callback:
            self.progress_callback(current_progress, max_progress)

    def get_settings(self):  # pragma: no cover
        """
        A parser must implement this
        """
        raise NotImplementedError

    def read_file_handle_unicode_errors(self, filepath: Path) -> str:
        """
        Helper utility for reading from a file, and handling a problem with its
        unicode, falling back to ignoring the error to remove the invalid bytes
        """
        try:
            text = filepath.read_text(encoding="utf-8")
        except UnicodeDecodeError as e:
            self.log.warning(f"Unicode error during text reading, continuing: {e}")
            text = filepath.read_bytes().decode("utf-8", errors="replace")
        return text

    def extract_metadata(self, document_path, mime_type):
        return []

    def get_page_count(self, document_path, mime_type) -> None:
        return None

    def parse(self, document_path, mime_type, file_name=None):
        raise NotImplementedError

    def get_archive_path(self):
        return self.archive_path

    def get_thumbnail(self, document_path, mime_type, file_name=None):
        """
        Returns the path to a file we can use as a thumbnail for this document.
        """
        raise NotImplementedError

    def get_text(self):
        return self.text

    def get_date(self) -> datetime.datetime | None:
        return self.date

    def cleanup(self) -> None:
        self.log.debug(f"Deleting directory {self.tempdir}")
        shutil.rmtree(self.tempdir)