anti-plagiarism/scripts/parsers/pmc_bulk.py

"""Парсер PubMed Central из открытого бакета AWS — массовый путь.

Отличие от `pmc.py`: тот ходит в API E-utilities и годится для точечной
докачки по теме (1-2 статьи в секунду, есть rate limit). Этот берёт статьи из
бакета `pmc-oa-opendata`, где у каждой лежит **уже извлечённый текст**, и
качает их пачками в несколько потоков — путь к сотням тысяч и миллионам.

Ключ для доступа не нужен, бакет открыт. Замер: 12 статей/с в 10 потоков.

Как и парсер Википедии, отдаёт статьи генератором: складывать миллионы
документов в список нельзя (см. bulk_writer.py).
"""

import logging
import re
from collections.abc import Iterator
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any
from urllib.parse import quote

import httpx
from base import BaseParser, ProgressCallback

logger = logging.getLogger(__name__)

BUCKET = "https://pmc-oa-opendata.s3.amazonaws.com"
KEY_RE = re.compile(r"<Key>(PMC\d+\.\d+)/\1\.txt</Key>")
TOKEN_RE = re.compile(r"<NextContinuationToken>([^<]+)</NextContinuationToken>")
YEAR_RE = re.compile(r"\b(19|20)\d{2}\b")
MIN_CHARS = 1500        # меньше — обрывок, а не статья


class PMCBulkParser(BaseParser):
    """PMC Open Access из бакета AWS. Отдаёт статьи потоком."""

    source_name = "pmc"     # тот же источник в корпусе, что и у API-парсера
    bulk = True

    def __init__(self, workers: int = 12) -> None:
        super().__init__()
        self.workers = workers
        # 60с на файл при бюджете задачи в 1500с — один залипший запрос съедал
        # почти весь бюджет. Обычный GET сюда укладывается в доли секунды,
        # 12с — с большим запасом на джиттер, но без риска съесть весь прогон
        self.client = httpx.Client(
            timeout=12,
            headers={"User-Agent": "AcademicHelper/1.0 (noreply@jze9.ru)"},
        )
        # Позиция листинга: бакет отдаётся страницами, и продолжать прогон
        # нужно с той же страницы, а не с начала
        self.last_token: str | None = None

    def fetch(  # type: ignore[override]
        self,
        limit: int = 1000,
        resume_token: str | None = None,
        start_after: str | None = None,
        progress_cb: ProgressCallback | None = None,
        **_ignored: Any,
    ) -> Iterator[dict[str, Any]]:
        """Статьи из бакета: листинг страницами, скачивание в потоках.

        Args:
            limit: сколько статей отдать
            resume_token: продолжить листинг с сохранённой страницы бакета
            start_after: ключ, после которого начинать листинг. Нужен, когда
                токена нет (первый прогон после ручных заливок): без него
                листинг пошёл бы с начала бакета и часами перемалывал уже
                залитые статьи как дубли
            progress_cb: см. base.ProgressCallback
        """
        token = resume_token
        given = 0

        while given < limit:
            try:
                ids, token = self._list_page(
                    token, want=min(200, limit - given),
                    start_after=None if token else start_after,
                )
            except Exception as e:
                logger.error("PMC bulk: листинг не удался: %s", e)
                return
            if not ids:
                logger.info("PMC bulk: бакет закончился")
                return

            # as_completed вместо map(): map() отдаёт результаты строго по
            # порядку отправки, поэтому один залипший запрос блокирует все
            # уже готовые — даже если остальные 11 потоков давно отработали
            with ThreadPoolExecutor(max_workers=self.workers) as pool:
                futures = {pool.submit(self._fetch_article, i): i for i in ids}
                for future in as_completed(futures):
                    raw = future.result()
                    if raw is None:
                        continue
                    given += 1
                    raw["resume_token"] = token
                    yield raw
                    if given >= limit:
                        break

            self.last_token = token
            if progress_cb and not progress_cb(given):
                logger.info("PMC bulk: выборка остановлена по запросу (%d)", given)
                return
            if token is None:
                return

    def _list_page(
        self, token: str | None, want: int, start_after: str | None = None
    ) -> tuple[list[str], str | None]:
        """Одна страница листинга бакета: id статей и токен следующей страницы."""
        url = f"{BUCKET}/?list-type=2&max-keys=1000"
        if token:
            # В токене бывают + и / — без экранирования S3 отвечает 400
            url += f"&continuation-token={quote(token, safe='')}"
        elif start_after:
            url += f"&start-after={quote(start_after, safe='')}"
        resp = self.client.get(url)
        resp.raise_for_status()
        ids = KEY_RE.findall(resp.text)[:want]
        m = TOKEN_RE.search(resp.text)
        return ids, (m.group(1) if m else None)

    def _fetch_article(self, art_id: str) -> dict[str, Any] | None:
        """Текст статьи и метаданные; None — статья недоступна или пустая."""
        try:
            txt = self.client.get(f"{BUCKET}/{art_id}/{art_id}.txt")
            if txt.status_code != 200 or len(txt.text) < MIN_CHARS:
                return None
            meta_resp = self.client.get(f"{BUCKET}/{art_id}/{art_id}.json")
            meta = meta_resp.json() if meta_resp.status_code == 200 else {}
        except Exception:
            return None

        return {"art_id": art_id, "meta": meta, "text": txt.text}

    def transform(self, raw: dict[str, Any]) -> dict[str, Any]:
        """Привести статью к унифицированному формату документов корпуса."""
        meta = raw.get("meta") or {}
        art_id = raw.get("art_id", "")
        pmcid = meta.get("pmcid") or art_id.split(".")[0]
        if not pmcid:
            return {}

        citation = meta.get("citation") or ""
        year_match = YEAR_RE.search(citation)
        text = raw.get("text", "")

        return {
            "source": self.source_name,
            "ext_id": f"pmc:{pmcid}",
            "title": (meta.get("title") or pmcid)[:1000],
            "authors": [],
            "doi": meta.get("doi"),
            "year": int(year_match.group(0)) if year_match else None,
            "lang": "en",
            "journal": citation[:300] or None,
            "url": f"https://www.ncbi.nlm.nih.gov/pmc/articles/{pmcid}/",
            "abstract": text[:2000],
            "text": text,
            "resume_token": raw.get("resume_token"),
        }