wiki-kb/wikikb/html2md.py

"""Convert MediaWiki `action=parse` HTML into clean markdown with frontmatter."""

from __future__ import annotations

import re

from bs4 import BeautifulSoup
from markdownify import markdownify as _markdownify

# Chrome / navigation / reference elements we don't want in the corpus.
_REMOVE_SELECTORS = (
    "style",
    "script",
    "link",
    ".mw-editsection",
    ".mw-jump-link",
    ".noprint",
    ".navbox",
    ".toc",
    "#toc",
    ".mw-empty-elt",
    ".reference",
    ".mw-references-wrap",
    ".reflist",
    ".ambox",
    ".metadata",
    ".pagetabber-absolute",
    ".pi-image",
)

_BLANK_LINES = re.compile(r"\n{3,}")


def _infobox_rows(element, group: str, tab: str, rows: list[tuple[str, str, str, str]]) -> None:
    """Collect (group, tab, label, value) from a portable-infobox subtree, in order.

    Fandom stat blocks are an <aside> of heading/value *siblings*, not a table, so
    markdownify alone emits a stream of headings followed by orphaned numbers — the
    values survive but stop being attached to their labels. This walk restores the
    pairing before any markdown is generated.
    """
    for child in element.find_all(recursive=False):
        classes = set(child.get("class") or [])
        if "pi-header" in classes:
            # A group heading scopes every following sibling until the next one.
            group = child.get_text(" ", strip=True)
        elif "wds-tabber" in classes:
            # Tier tabs (I/II/III/IV) are all inlined; without their label the four
            # stat blocks are indistinguishable.
            labels = [t.get_text(" ", strip=True) for t in child.select(".wds-tabs__tab-label")]
            for index, panel in enumerate(child.select(".wds-tab__content")):
                _infobox_rows(panel, group, labels[index] if index < len(labels) else "", rows)
        elif "pi-smart-group" in classes:
            # Horizontal group: labels in one row, values positionally in the next.
            labels = [h.get_text(" ", strip=True) for h in child.select(".pi-smart-data-label")]
            values = [v.get_text(" ", strip=True) for v in child.select(".pi-smart-data-value")]
            # strict=False: a hand-edited infobox with a stray label and no value
            # should degrade to the pairs that do line up, not abort the scrape.
            rows.extend(
                (group, tab, label, value) for label, value in zip(labels, values, strict=False)
            )
        elif "pi-data" in classes:
            label = child.select_one(".pi-data-label")
            value = child.select_one(".pi-data-value")
            if label is not None and value is not None:
                rows.append(
                    (group, tab, label.get_text(" ", strip=True), value.get_text(" ", strip=True))
                )
        else:
            _infobox_rows(child, group, tab, rows)


def _rows_to_html(rows: list[tuple[str, str, str, str]]) -> str:
    """Render collected infobox rows as heading + table HTML, grouped in order."""
    parts: list[str] = []
    current: tuple[str, str] | None = None
    for group, tab, label, value in rows:
        if (group, tab) != current:
            if current is not None:
                parts.append("</table>")
            current = (group, tab)
            heading = "Infobox" if not group else f"Infobox: {group}"
            if tab:
                heading = f"{heading} ({tab})"
            parts.append(f"<h2>{heading}</h2>")
            parts.append("<table><tr><th>Field</th><th>Value</th></tr>")
        parts.append(f"<tr><td>{label}</td><td>{value}</td></tr>")
    if current is not None:
        parts.append("</table>")
        # The infobox emits headings, so any lead paragraph that follows it would
        # otherwise be filed under the LAST infobox group ("Infobox: Technical") —
        # mislabelling the page's defining sentence, which is the worst possible
        # chunk to misname given section text is weighted in both BM25 and the
        # embedding. Give that text its own section instead; if a real heading
        # follows immediately, this one holds no text and the chunker drops it.
        parts.append("<h2>Overview</h2>")
    return "".join(parts)


def transform_infoboxes(root) -> None:
    """Replace every portable-infobox in `root` with label/value tables, in place."""
    for aside in root.select("aside.portable-infobox"):
        rows: list[tuple[str, str, str, str]] = []
        _infobox_rows(aside, "", "", rows)
        replacement = _rows_to_html(rows)
        if replacement:
            aside.replace_with(BeautifulSoup(replacement, "html.parser"))
        else:
            # Theme/dummy infoboxes carry no data; dropping them avoids an empty
            # zero-width-space heading in the corpus.
            aside.decompose()


def clean_html(html: str) -> str:
    """Strip nav/edit/reference chrome, returning the article body HTML."""
    soup = BeautifulSoup(html, "html.parser")
    root = soup.select_one(".mw-parser-output") or soup
    for selector in _REMOVE_SELECTORS:
        for element in root.select(selector):
            element.decompose()
    transform_infoboxes(root)
    return str(root)


def html_to_markdown(html: str) -> str:
    """Rendered HTML fragment -> collapsed ATX markdown (images dropped)."""
    cleaned = clean_html(html)
    markdown = _markdownify(cleaned, heading_style="ATX", strip=["img"])
    return _BLANK_LINES.sub("\n\n", markdown).strip()


def _yaml_scalar(value: str) -> str:
    """Quote a string so it is a safe single-line YAML scalar."""
    escaped = value.replace("\\", "\\\\").replace('"', '\\"')
    return f'"{escaped}"'


def build_document(
    *,
    title: str,
    url: str,
    revid: int,
    categories: list[str],
    html: str,
    source: str = "mediawiki",
    variant: str = "unknown",
) -> str:
    """Assemble a corpus markdown document with YAML frontmatter."""
    body = html_to_markdown(html)
    category_list = ", ".join(_yaml_scalar(c) for c in categories)
    return "\n".join(
        (
            "---",
            f"title: {_yaml_scalar(title)}",
            f"url: {url}",
            f"revid: {revid}",
            f"source: {source}",
            f"variant: {variant}",
            f"categories: [{category_list}]",
            "---",
            "",
            f"# {title}",
            "",
            body,
            "",
        )
    )