wiki-kb/wikikb/sources.py

"""Pluggable page sources — where a wiki's pages come from.

Everything downstream of "fetch -> HTML" is already source-agnostic: html2md,
chunker, embed, indexer, search and server never learn where the HTML originated.
The wiki-engine assumption is confined to exactly two operations, which is what this
module abstracts:

  * **enumerate** the content pages, and
  * **fetch** one page as server-rendered HTML.

Two implementations ship:

  * :class:`MediaWikiSource` — the good case. An open ``api.php`` gives a real page
    list, server-rendered HTML with templates expanded, categories, revision ids for
    fidelity diffing, and ``maxlag`` so the server can push back.
  * :class:`SitemapHtmlSource` — the fallback for wikis with no API at all (custom
    engines such as Fextralife). Enumerates ``sitemap.xml`` and extracts a content
    container from each page. Works, but gives up revision ids and ``maxlag``.

Use :func:`detect_source` to pick one by probing the host.
"""

from __future__ import annotations

import re
import urllib.parse
from collections.abc import Callable, Iterator
from dataclasses import dataclass, field
from typing import Protocol

import httpx
from bs4 import BeautifulSoup

from . import config

MEDIAWIKI = "mediawiki"
HTML = "html"

# Candidate content containers for HTML wikis, most specific first. Fextralife uses
# #main-content; MediaWiki mirrors use .mw-parser-output; the rest are conventional.
_CONTENT_SELECTORS = (
    "#main-content",
    "#wiki-content-block",
    ".mw-parser-output",
    "#content",
    "main",
    "article",
)
_LOC = re.compile(r"<loc>\s*([^<\s]+)\s*</loc>", re.IGNORECASE)


@dataclass(frozen=True)
class PageDoc:
    """A fetched page ready to be written to the corpus."""

    title: str
    url: str
    html: str
    source: str
    # Revision id, or 0 for sources that have no revision concept. A real revid lets
    # the fidelity check say "the wiki was edited" rather than "our extractor broke";
    # without one, only a content diff is available.
    revid: int = 0
    categories: list[str] = field(default_factory=list)


class Source(Protocol):
    """Where pages come from. Implementations must be polite and resumable."""

    name: str

    def iter_titles(self, *, limit: int | None = None) -> Iterator[str]:
        """Yield content-page titles."""
        ...

    def fetch_page(self, title: str) -> PageDoc | None:
        """Fetch one page, or None if it is missing."""
        ...


class MediaWikiSource:
    """Pages from a MediaWiki ``api.php``: page list + server-rendered HTML."""

    name = MEDIAWIKI

    def __init__(self, api_get: Callable[..., dict], *, api_url: str | None = None) -> None:
        self._api_get = api_get
        self._api_url = api_url or config.API_URL

    def iter_titles(self, *, limit: int | None = None) -> Iterator[str]:
        apcontinue: str | None = None
        seen = 0
        while True:
            params: dict[str, object] = {
                "action": "query",
                "list": "allpages",
                "apnamespace": config.CONTENT_NAMESPACE,
                "apfilterredir": "nonredirects",
                "aplimit": "500",
            }
            if apcontinue is not None:
                params["apcontinue"] = apcontinue
            data = self._api_get(params)
            for page in data["query"]["allpages"]:
                yield page["title"]
                seen += 1
                if limit is not None and seen >= limit:
                    return
            cont = data.get("continue")
            if not cont:
                return
            apcontinue = cont["apcontinue"]

    def fetch_page(self, title: str) -> PageDoc | None:
        data = self._api_get(
            {
                "action": "parse",
                "page": title,
                "prop": "text|categories|revid",
                "redirects": "1",
            }
        )
        if "error" in data:
            return None
        parse = data["parse"]
        return PageDoc(
            title=parse["title"],
            url=config.page_url(parse["title"]),
            html=parse["text"],
            source=self.name,
            revid=parse["revid"],
            categories=[c["category"].replace("_", " ") for c in parse.get("categories", [])],
        )


class SitemapHtmlSource:
    """Pages from ``sitemap.xml`` + a content container in each page's HTML.

    For wikis with no API. Two capabilities are genuinely lost and callers should not
    pretend otherwise: there is no revision id (``revid`` stays 0), and there is no
    ``maxlag``, so the rate limit is purely self-imposed with no server backpressure.
    """

    name = HTML

    def __init__(
        self,
        get: Callable[[str], httpx.Response],
        *,
        sitemap_url: str | None = None,
        selectors: tuple[str, ...] = _CONTENT_SELECTORS,
    ) -> None:
        self._get = get
        self._sitemap_url = sitemap_url or f"{config.WIKI_HOST}/sitemap.xml"
        self._selectors = selectors
        self._urls: dict[str, str] = {}

    @staticmethod
    def title_from_url(url: str) -> str:
        """Last path segment as a human title ("/Common+Longsword" -> "Common Longsword")."""
        slug = urllib.parse.urlparse(url).path.rstrip("/").rsplit("/", 1)[-1]
        return urllib.parse.unquote_plus(slug).replace("_", " ").strip()

    def iter_titles(self, *, limit: int | None = None) -> Iterator[str]:
        text = self._get(self._sitemap_url).text
        locs = _LOC.findall(text)
        # A sitemap index points at further sitemaps rather than pages.
        if locs and all(loc.endswith((".xml", ".xml.gz")) for loc in locs):
            nested: list[str] = []
            for child in locs:
                nested.extend(_LOC.findall(self._get(child).text))
            locs = nested
        seen = 0
        for url in locs:
            title = self.title_from_url(url)
            if not title:
                continue
            self._urls.setdefault(title, url)
            yield title
            seen += 1
            if limit is not None and seen >= limit:
                return

    def _extract(self, html: str) -> str | None:
        soup = BeautifulSoup(html, "html.parser")
        for selector in self._selectors:
            node = soup.select_one(selector)
            if node is not None:
                return str(node)
        return None

    def fetch_page(self, title: str) -> PageDoc | None:
        url = self._urls.get(title) or f"{config.WIKI_HOST}/{urllib.parse.quote(title)}"
        response = self._get(url)
        if response.status_code == 404:
            return None
        body = self._extract(response.text)
        if body is None:
            return None
        return PageDoc(title=title, url=url, html=body, source=self.name)


def detect_source_kind(get: Callable[[str], httpx.Response], host: str) -> tuple[str, str]:
    """Probe a host and return (kind, endpoint).

    Order matters: an open MediaWiki API beats HTML scraping every time, so try both
    conventional API paths ("/api.php" on Fandom, wiki.gg and paradoxwikis;
    "/w/api.php" on Wikipedia) before falling back to the sitemap. Raises if neither
    is available — that is a real answer, not a failure to be papered over.
    """
    for path in ("/api.php", "/w/api.php"):
        endpoint = f"{host}{path}"
        try:
            response = get(f"{endpoint}?action=query&meta=siteinfo&format=json")
            if response.status_code == 200 and "mediawiki" in (
                response.json().get("query", {}).get("general", {}).get("generator", "").lower()
            ):
                return MEDIAWIKI, endpoint
        except (httpx.HTTPError, ValueError, KeyError):
            continue
    sitemap = f"{host}/sitemap.xml"
    response = get(sitemap)
    if response.status_code == 200 and "<loc>" in response.text.lower():
        return HTML, sitemap
    raise RuntimeError(
        f"{host} exposes neither a MediaWiki API nor a sitemap.xml; "
        "it cannot be mirrored by this tool."
    )