Literature Search

#!/usr/bin/env python3
"""Download arXiv paper source by title, extract .tex content.

Searches the arXiv API by title, downloads the source tarball,
and extracts .tex files into a local directory.

Self-contained: uses only stdlib (urllib + xml.etree instead of feedparser).

Usage:
    python download_arxiv_source.py --title "Attention Is All You Need" --output-dir arxiv_papers/
    python download_arxiv_source.py --title "BERT" --max-results 3 --output-dir arxiv_papers/
    python download_arxiv_source.py --arxiv-id 1706.03762 --output-dir arxiv_papers/
"""

import argparse
import json
import os
import re
import sys
import tarfile
import tempfile
import time
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET

ARXIV_API = "http://export.arxiv.org/api/query"
ARXIV_NS = {"atom": "http://www.w3.org/2005/Atom"}


def search_arxiv(query: str, max_results: int = 5, search_field: str = "ti") -> list[dict]:
    """Search arXiv API and return paper metadata."""
    params = urllib.parse.urlencode({
        "search_query": f"{search_field}:{urllib.parse.quote(query)}",
        "start": 0,
        "max_results": max_results,
        "sortBy": "relevance",
        "sortOrder": "descending",
    })
    url = f"{ARXIV_API}?{params}"

    try:
        with urllib.request.urlopen(url, timeout=30) as resp:
            xml_data = resp.read()
    except Exception as e:
        print(f"arXiv API error: {e}", file=sys.stderr)
        return []

    root = ET.fromstring(xml_data)
    papers = []
    for entry in root.findall("atom:entry", ARXIV_NS):
        title_el = entry.find("atom:title", ARXIV_NS)
        title = title_el.text.strip().replace("\n", " ") if title_el is not None else ""

        summary_el = entry.find("atom:summary", ARXIV_NS)
        summary = summary_el.text.strip() if summary_el is not None else ""

        authors = []
        for author in entry.findall("atom:author", ARXIV_NS):
            name_el = author.find("atom:name", ARXIV_NS)
            if name_el is not None:
                authors.append(name_el.text)

        published_el = entry.find("atom:published", ARXIV_NS)
        published = published_el.text if published_el is not None else ""

        abs_link = ""
        pdf_link = ""
        for link in entry.findall("atom:link", ARXIV_NS):
            href = link.get("href", "")
            link_type = link.get("type", "")
            rel = link.get("rel", "")
            if link_type == "application/pdf":
                pdf_link = href
            elif rel == "alternate":
                abs_link = href

        arxiv_id = ""
        id_el = entry.find("atom:id", ARXIV_NS)
        if id_el is not None:
            m = re.search(r"abs/(.+)", id_el.text)
            if m:
                arxiv_id = m.group(1)

        papers.append({
            "title": title,
            "authors": authors,
            "published": published,
            "summary": summary,
            "abs_link": abs_link,
            "pdf_link": pdf_link,
            "arxiv_id": arxiv_id,
        })

    return papers


def download_source(arxiv_id: str, output_dir: str) -> str | None:
    """Download arXiv source tarball and extract .tex files.

    Returns the path to the extracted content or None on failure.
    """
    # Strip version suffix for source download
    base_id = re.sub(r"v\d+$", "", arxiv_id)
    source_url = f"https://arxiv.org/src/{base_id}"

    safe_name = re.sub(r"[^a-zA-Z0-9._-]", "_", arxiv_id)
    os.makedirs(output_dir, exist_ok=True)

    try:
        req = urllib.request.Request(source_url, headers={"User-Agent": "SkillScript/1.0"})
        with urllib.request.urlopen(req, timeout=60) as resp:
            tar_data = resp.read()
    except Exception as e:
        print(f"Download failed for {arxiv_id}: {e}", file=sys.stderr)
        return None

    # Save tarball to temp file, then extract
    with tempfile.NamedTemporaryFile(suffix=".tar.gz", delete=False) as tmp:
        tmp.write(tar_d