signal-bot/tool/commons_harvest.py

#!/usr/bin/env python3
# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence.
"""Harvesting files out of Commons categories, licence and all.

Split from ``commons_api`` at the 250-line cap. That module is the transport
-- one request, retried politely -- and this one is everything built on it:
which categories a pool draws from, which files may be reposted, and the
cache that lets an interrupted harvest resume instead of starting over.
"""

from __future__ import annotations

import sys
from concurrent.futures import ThreadPoolExecutor
from functools import partial
from pathlib import Path
from typing import Any

from commons_api import (
    BATCH,
    LICENCES,
    SUFFIXES,
    WORKERS,
    get,
    licence_of,
    members,
)


def licensed_members(category: str, *, user_agent: str) -> dict[int, str]:
    """Every usable, repostable file in one category, in a single pass.

    ``list=categorymembers`` names the files and ``prop=imageinfo`` reports
    the licence, and MediaWiki will do both in one request if the first is
    run as a *generator* for the second. Doing them separately -- which is
    what this module did first, and what gen_traps still does -- costs an
    extra request per fifty files, and those requests are the slow kind:
    measured at 2.3s each against 0.75s for a plain listing, because
    extmetadata is expensive. On a 75,000-file harvest that second pass was
    forty-six minutes on its own.

    Verified equivalent to ``members`` + ``licensed`` page id for page id, on
    a painter's category of 387 files and on Category:Ukiyo-e of 311.
    """
    kept: dict[int, str] = {}
    carry: dict[str, str] = {}
    while True:
        document = get(
            {
                "generator": "categorymembers",
                "gcmtitle": category,
                "gcmtype": "file",
                "gcmlimit": "500",
                "prop": "imageinfo",
                "iiprop": "extmetadata",
                "iiextmetadatafilter": "LicenseShortName",
                **carry,
            },
            user_agent,
        )
        for page in document.get("query", {}).get("pages", []):
            title = str(page["title"])
            if Path(title).suffix.lower() not in SUFFIXES:
                continue
            if licence_of(page).casefold().startswith(LICENCES):
                kept[int(page["pageid"])] = title
        carry = document.get("continue", {})
        if not carry:
            return kept


def _fetch_licensed(
    wanted: list[str],
    user_agent: str,
    cache: dict[str, dict[int, str]] | None,
) -> dict[int, str]:
    """Every licensed file in these categories, reusing whatever is cached.

    Split out of walk because it is the only part that is concurrent and the
    only part that touches the cache; walk stays a description of which
    categories get visited, which is what makes the pool ordering readable.
    """
    files: dict[int, str] = {}
    if cache is not None:
        for category in [name for name in wanted if name in cache]:
            files.update(cache[category])
        wanted = [name for name in wanted if name not in cache]
    # A quotation mark would end the incategory: term early and silently
    # change the query, so those few categories keep the per-category path.
    # Per category, via list=categorymembers, because that is the
    # authoritative enumeration. Asking CirrusSearch instead (incategory:
    # ORed across a batch) looked like a 6x win and was measured to be
    # neither: it returned 2404 files where this returns 2429, dropped a
    # different ~1% on every run, and was slower once the overflow path
    # started firing. A search index is ranked and eventually consistent;
    # deep pagination over one silently loses rows. Kept as a comment so the
    # next person does not spend an afternoon rediscovering it.
    with ThreadPoolExecutor(max_workers=WORKERS) as pool:
        found_each = pool.map(partial(licensed_members, user_agent=user_agent), wanted)
        for category, found in zip(wanted, found_each, strict=True):
            files.update(found)
            if cache is not None:
                cache[category] = found
    return files


def walk(
    seeds: tuple[str, ...],
    excluded: frozenset[str],
    patterns: tuple[str, ...],
    user_agent: str,
    seen: set[str] | None = None,
    *,
    with_licences: bool = False,
    cache: dict[str, dict[int, str]] | None = None,
) -> dict[int, str]:
    """Page id to title, for every usable file one step from the seeds.

    Depth is one, and that is measured rather than chosen: recursing three
    levels from *Category:Fish traps* reaches Alfred Sisley's net paintings
    and Hawaiian fisherman watercolours. Immediate subcategories only.

    Args:
        seeds: Categories to start from.
        excluded: Categories pruned as whole subtrees.
        patterns: A subcategory whose title contains one of these is pruned.
        user_agent: Passed through to every request.
        seen: Categories already walked. Pass one set across two calls to keep
            two pools disjoint -- a file in both would be drawn at both odds.
        with_licences: Filter licences in the same request that lists the
            files, rather than in a second pass. Faster by far -- see
            licensed_members -- and the caller then must not call
            ``licensed`` again.
        cache: Category to its files, from load_cache. Anything already in
            here is not fetched again, which is what makes an interrupted
            harvest resumable instead of wasted. Mutated as the walk goes.

    Returns:
        Every file found, page id to title.
    """
    files: dict[int, str] = {}
    visited = set(excluded) if seen is None else seen | set(excluded)
    for number, seed in enumerate(seeds, start=1):
        sys.stderr.write(f"[{number}/{len(seeds)}] {seed} ({len(files)} files so far)\n")
        sys.stderr.flush()
        branch = [seed, *(sub["title"] for sub in members(seed, "subcat", user_agent))]
        wanted: list[str] = []
        for category in branch:
            if category in visited or any(pattern in category for pattern in patterns):
                continue
            visited.add(category)
            wanted.append(category)
        if with_licences:
            files.update(_fetch_licensed(wanted, user_agent, cache))
            continue
        for category in wanted:
            for entry in members(category, "file", user_agent):
                title = str(entry["title"])
                if Path(title).suffix.lower() in SUFFIXES:
                    files[int(entry["pageid"])] = title
    if seen is not None:
        seen |= visited
    return files


def _titles_page(titles: list[str], user_agent: str) -> dict[str, Any]:
    """One imageinfo lookup for a batch of file titles."""
    return get(
        {
            "titles": "|".join(titles),
            "prop": "imageinfo",
            "iiprop": "extmetadata",
            "iiextmetadatafilter": "LicenseShortName",
        },
        user_agent,
    )


def licensed_titles(titles: list[str], user_agent: str) -> dict[int, str]:
    """Resolve Commons file titles to page ids, keeping only repostable ones.

    The mirror image of ``licensed``: that one starts from page ids the walk
    found, this one starts from titles somebody else named -- Wikidata, which
    knows the filename and not the id. Same licence rule, same batch size.
    """
    kept: dict[int, str] = {}
    batches = [titles[start : start + BATCH] for start in range(0, len(titles), BATCH)]
    with ThreadPoolExecutor(max_workers=WORKERS) as pool:
        documents = list(pool.map(partial(_titles_page, user_agent=user_agent), batches))
    for document in documents:
        for page in document.get("query", {}).get("pages", []):
            if page.get("missing"):
                continue
            title = str(page["title"])
            if Path(title).suffix.lower() not in SUFFIXES:
                continue
            if licence_of(page).casefold().startswith(LICENCES):
                kept[int(page["pageid"])] = title
    return kept


def licensed(files: dict[int, str], user_agent: str) -> dict[int, str]:
    """Drop everything whose licence does not permit reposting it."""
    kept: dict[int, str] = {}
    ids = sorted(files)
    sys.stderr.write(f"checking licences on {len(ids)} files\n")
    sys.stderr.flush()
    for start in range(0, len(ids), BATCH):
        chunk = ids[start : start + BATCH]
        document = get(
            {
                "prop": "imageinfo",
                "iiprop": "extmetadata",
                "iiextmetadatafilter": "LicenseShortName",
                "pageids": "|".join(str(page) for page in chunk),
            },
            user_agent,
        )
        for page in document.get("query", {}).get("pages", []):
            if licence_of(page).casefold().startswith(LICENCES):
                kept[int(page["pageid"])] = files[int(page["pageid"])]
    return kept