html-to-pdf/app/pdf/chunker.py

"""Splitting of a large document into renderable chunks.

A thousand page document rendered in one pass keeps the whole page tree in
memory. Splitting it on structural boundaries keeps memory flat at the cost of
having to reassemble the page numbering afterwards.

Cuts are only ever made between direct children of the container element, so a
table or a paragraph is never torn in half.
"""

from __future__ import annotations

import logging

from lxml import html as lxml_html

from .document import SourceDocument

logger = logging.getLogger(__name__)

SPLIT_TAGS = {"section", "article", "h1"}
BREAK_KEYWORDS = ("page-break-before", "break-before")

# Rough page size heuristic. The real page count is only known after rendering,
# so this only has to be good enough to keep chunks roughly even.
CHARS_PER_PAGE = 2200
IMAGE_CHAR_WEIGHT = 900
TABLE_ROW_CHAR_WEIGHT = 120


def is_split_point(element) -> bool:
    if not isinstance(element.tag, str):
        return False
    if element.tag.lower() in SPLIT_TAGS:
        return True
    if element.get("data-chunk") is not None:
        return True
    style = (element.get("style") or "").lower()
    return any(keyword in style for keyword in BREAK_KEYWORDS)


def estimate_pages(element) -> float:
    text_length = len(element.text_content())
    text_length += IMAGE_CHAR_WEIGHT * len(element.findall(".//img"))
    text_length += TABLE_ROW_CHAR_WEIGHT * len(element.findall(".//tr"))
    return max(text_length / CHARS_PER_PAGE, 0.01)


def split_document(document: SourceDocument, pages_per_chunk: int) -> list[str]:
    """Return one complete HTML document per chunk.

    Falls back to a single chunk when the document has no usable split points.
    """
    container = document.container
    children = [child for child in container if isinstance(child.tag, str)]

    split_indexes = [index for index, child in enumerate(children) if is_split_point(child)]
    if len(split_indexes) < 2:
        logger.info(
            "Document has no usable split points, rendering in one pass",
            extra={"split_points": len(split_indexes), "children": len(children)},
        )
        return [document.to_html()]

    groups = _group_children(children, set(split_indexes), pages_per_chunk)
    if len(groups) < 2:
        logger.info("Document fits into a single chunk", extra={"children": len(children)})
        return [document.to_html()]

    prefix, suffix = _skeleton(document, container)
    child_html = [lxml_html.tostring(child, encoding="unicode") for child in children]

    chunks = ["".join((prefix, *(child_html[index] for index in group), suffix)) for group in groups]
    logger.info(
        "Document split into chunks",
        extra={"chunks": len(chunks), "children": len(children), "pages_per_chunk": pages_per_chunk},
    )
    return chunks


def _group_children(children, split_indexes: set[int], pages_per_chunk: int) -> list[list[int]]:
    groups: list[list[int]] = []
    current: list[int] = []
    current_pages = 0.0

    for index, child in enumerate(children):
        starts_chunk = index in split_indexes and current and current_pages >= pages_per_chunk
        if starts_chunk:
            groups.append(current)
            current = []
            current_pages = 0.0

        current.append(index)
        current_pages += estimate_pages(child)

    if current:
        groups.append(current)

    return groups


def _skeleton(document: SourceDocument, container) -> tuple[str, str]:
    """Opening and closing markup shared by every chunk.

    The whole ancestor chain is recreated with its attributes so CSS selectors
    that depend on it keep matching inside a chunk.
    """
    head_html = lxml_html.tostring(document.head, encoding="unicode")

    chain = []
    node = container
    while node is not None and node is not document.tree:
        chain.append(node)
        node = node.getparent()
    chain.reverse()

    opens = [_open_tag(document.tree)]
    opens.append(head_html)
    closes = ["</html>"]

    for node in chain:
        opens.append(_open_tag(node))
        closes.append(f"</{node.tag}>")

    closes.reverse()
    return "<!DOCTYPE html>\n" + "".join(opens), "".join(closes)


def _open_tag(element) -> str:
    attributes = "".join(
        f' {name}="{_escape(value)}"' for name, value in element.attrib.items()
    )
    return f"<{element.tag}{attributes}>"


def _escape(value: str) -> str:
    return (
        value.replace("&", "&amp;")
        .replace('"', "&quot;")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
    )