codex-py/codex/sources/arxiv.py

"""arXiv API client.

Provides:
- fetch_metadata: resolve an arXiv ID to a Paper via the Atom export API.
- fetch_source: download the .tar.gz source of a paper and extract the primary .tex file.
- fetch_pdf_url: return the canonical PDF URL for a given arXiv ID.
"""

from __future__ import annotations

import gzip
import io
import logging
import tarfile
import xml.etree.ElementTree as ET

import httpx
from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential

from codex.models import Paper

logger = logging.getLogger(__name__)

_BASE = "https://arxiv.org"
_EXPORT = "https://export.arxiv.org"
_ATOM = "{http://www.w3.org/2005/Atom}"


def _is_retryable(exc: BaseException) -> bool:
    if isinstance(exc, httpx.HTTPStatusError):
        return exc.response.status_code == 429 or exc.response.status_code >= 500
    return False


@retry(
    retry=retry_if_exception(_is_retryable),
    stop=stop_after_attempt(4),
    wait=wait_exponential(min=1, max=20),
    before_sleep=lambda rs: logger.warning("arXiv retry %d", rs.attempt_number),
)
def _query(arxiv_id: str) -> httpx.Response:
    response = httpx.get(
        f"{_EXPORT}/api/query",
        params={"id_list": arxiv_id},
        timeout=30,
        follow_redirects=True,
    )
    response.raise_for_status()
    return response


def fetch_metadata(arxiv_id: str) -> Paper | None:
    """Resolve an arXiv ID to a Paper via the Atom export API.

    Authoritative metadata source for arXiv preprints — used as an ingest
    fallback when OpenAlex 404s on an arXiv id (DQ-2). The arXiv id (bare,
    e.g. ``"1911.00966"`` or legacy ``"math/0603097"``) becomes ``Paper.id``;
    ``openalex_id`` and ``bibkey`` are left for the caller to populate.

    Returns
    -------
    Paper | None
        Populated Paper (title, authors, year, abstract), or None if arXiv
        has no entry for the id.
    """
    bare = arxiv_id[len("arxiv:") :] if arxiv_id.lower().startswith("arxiv:") else arxiv_id
    try:
        response = _query(bare)
    except httpx.HTTPError:
        logger.warning("arXiv metadata fetch failed for %s", bare, exc_info=True)
        return None
    try:
        root = ET.fromstring(response.text)
    except ET.ParseError:
        return None
    entry = root.find(f"{_ATOM}entry")
    if entry is None:
        return None
    # An id-not-found query still returns a feed but with no <entry>.
    title = (entry.findtext(f"{_ATOM}title") or "").strip()
    if not title:
        return None
    summary = (entry.findtext(f"{_ATOM}summary") or "").strip()
    published = entry.findtext(f"{_ATOM}published") or ""
    year = int(published[:4]) if published[:4].isdigit() else None
    authors = [
        name.strip()
        for a in entry.findall(f"{_ATOM}author")
        if (name := a.findtext(f"{_ATOM}name")) and name.strip()
    ]
    return Paper(
        id=bare,
        title=" ".join(title.split()),
        authors=authors,
        year=year,
        abstract=" ".join(summary.split()) or None,
    )


def _has_documentclass(content: str) -> bool:
    """True if *content* has an UN-commented ``\\documentclass`` line.

    A plain ``"\\documentclass" in content`` substring test also matches a
    commented-out ``%\\documentclass`` (common in arXiv preambles); requiring the
    command to start its line avoids selecting the wrong primary file.
    """
    return any(line.lstrip().startswith("\\documentclass") for line in content.splitlines())


def fetch_source(arxiv_id: str) -> str | None:
    """Download and extract the primary LaTeX source for an arXiv paper.

    Downloads the .tar.gz source bundle from ``https://arxiv.org/src/{arxiv_id}``,
    locates the primary .tex file (preferring any file containing ``\\documentclass``,
    falling back to the largest .tex by size), and inlines its ``\\input``/
    ``\\include`` directives from the other archive members so multi-file projects
    return the full document body, not just the primary file's include skeleton.

    Parameters
    ----------
    arxiv_id:
        The arXiv identifier (e.g. ``"2301.07041"``).

    Returns
    -------
    str | None
        Raw LaTeX source string, or None if the paper is not found or no .tex
        file is present (signals Nougat fallback).
    """
    url = f"{_BASE}/src/{arxiv_id}"
    response = httpx.get(url, timeout=60, follow_redirects=True)
    if response.status_code == 404:
        logger.debug("arXiv 404 for source id=%s", arxiv_id)
        return None
    if response.status_code != 200:
        response.raise_for_status()

    from codex.parsing.tex import flatten_inputs

    raw = response.content
    # Image/font/binary members are never \input'd as text; skip them so the source
    # map stays text-only (decode-with-replace would otherwise store garbled bytes).
    skip_ext = (".png", ".jpg", ".jpeg", ".gif", ".bmp", ".pdf", ".eps", ".ps", ".ttf", ".otf")
    try:
        with tarfile.open(fileobj=io.BytesIO(raw), mode="r:gz") as tf:
            files: dict[str, str] = {}
            primary_name: str | None = None
            for member in tf.getmembers():
                if not member.isfile() or member.name.lower().endswith(skip_ext):
                    continue
                f = tf.extractfile(member)
                if f is None:
                    continue
                # Keep ALL text members resolvable so \input of a non-.tex include
                # (e.g. a .def or extension-less body file) is not silently dropped (R-C).
                files[member.name] = f.read().decode("utf-8", errors="replace")
                # Primary doc = a .tex with an UN-commented \documentclass.
                if (
                    primary_name is None
                    and member.name.endswith(".tex")
                    and _has_documentclass(files[member.name])
                ):
                    primary_name = member.name

            tex_files = {k: v for k, v in files.items() if k.endswith(".tex")}
            if not tex_files:
                logger.debug("No .tex files found in arXiv source for %s", arxiv_id)
                return None

            # No \documentclass anywhere → fall back to the largest .tex.
            if primary_name is None:
                primary_name = max(tex_files, key=lambda k: len(tex_files[k]))

            # Inline \input/\include so multi-file projects yield the full body,
            # not just the primary file's skeleton of include directives (R-C).
            return flatten_inputs(files[primary_name], files)
    except tarfile.TarError:
        # Old arXiv submissions are served as a single bare gzipped .tex with no
        # tar wrapper (e.g. legacy math/* papers) — gunzip the body directly.
        try:
            single = gzip.decompress(raw).decode("utf-8", errors="replace")
        except (OSError, EOFError) as exc:
            logger.warning("Failed to read arXiv source for %s: %s", arxiv_id, exc)
            return None
        if "\\" in single[:5000]:  # looks like LaTeX (has backslash commands)
            return single
        logger.debug("arXiv source for %s is not LaTeX", arxiv_id)
        return None


def fetch_pdf_url(arxiv_id: str) -> str:
    """Return the canonical PDF URL for an arXiv paper.

    This is a pure computation — no HTTP request is made.

    Parameters
    ----------
    arxiv_id:
        The arXiv identifier (e.g. ``"2301.07041"``).

    Returns
    -------
    str
        The full URL of the PDF file.
    """
    return f"{_BASE}/pdf/{arxiv_id}.pdf"