jcode/skills/scientific-deep-research/papers.py

367 lines
18 KiB
Python

#!/usr/bin/env python3
"""papers.py - find, resolve, download and cite scientific papers (stdlib only).
python papers.py search "boolean operations on triangle meshes" --max 25 [--since 2018]
python papers.py resolve 1706.03762 # arXiv id, DOI or "title": find the peer-reviewed version
python papers.py bibtex 10.1145/3306346.3322959 # DOI / arXiv id / dblp key -> BibTeX entry
python papers.py download 1706.03762 --out papers/ # PDF (arXiv or open-access copy)
python papers.py text papers/vaswani2017attention.pdf [--pages 12] # PDF -> .txt (pymupdf)
search merges arXiv, OpenAlex and Semantic Scholar; prints a ranked table and writes candidates.json.
resolve prints, for one paper, the best version: journal/conference (DOI, venue, year) if one
exists, otherwise the arXiv preprint - plus the BibTeX key you should use.
bibtex prints a cleaned BibTeX entry with a stable key <lastname><year><firstword>.
All network calls have a 20 s timeout and identify themselves; be polite (<= 1 request/s).
Sources: export.arxiv.org (Atom), api.openalex.org, api.semanticscholar.org, api.crossref.org,
dblp.org, doi.org (content negotiation). Set CONTACT_EMAIL for the Crossref/OpenAlex polite pools.
"""
import json, os, re, sys, time, unicodedata, urllib.parse, urllib.request
import xml.etree.ElementTree as ET
from pathlib import Path
UA = f"papers.py/1.0 (mailto:{os.environ.get('CONTACT_EMAIL', 'research@example.org')})"
ARXIV_RE = re.compile(r"(?:arxiv[:/]|abs/|pdf/)?(\d{4}\.\d{4,5})(?:v\d+)?", re.I)
DOI_RE = re.compile(r"\b(10\.\d{4,9}/[^\s\"<>]+)", re.I)
# ----------------------------------------------------------------- helpers
def get(url, accept="application/json", retries=2):
req = urllib.request.Request(url, headers={"User-Agent": UA, "Accept": accept})
for attempt in range(retries + 1):
try:
with urllib.request.urlopen(req, timeout=20) as r:
return r.read().decode("utf-8", "replace")
except Exception as e: # noqa: BLE001
if attempt == retries:
print(f" ! {url[:80]}...: {e}", file=sys.stderr)
return None
time.sleep(2 * (attempt + 1))
def norm_title(t):
t = unicodedata.normalize("NFKD", t or "").encode("ascii", "ignore").decode()
return re.sub(r"[^a-z0-9]+", " ", t.lower()).strip()
def bib_key(authors, year, title):
last = (authors[0].split()[-1] if authors else "anon")
last = re.sub(r"[^a-z]", "", unicodedata.normalize("NFKD", last).encode("ascii", "ignore").decode().lower()) or "anon"
stop = {"a", "an", "the", "on", "of", "for", "and", "to", "in", "with", "towards", "toward"}
first = next((w for w in re.findall(r"[a-z0-9]+", (title or "").lower()) if w not in stop), "paper")
return f"{last}{year or 'nd'}{first}"
def is_arxiv_doi(doi):
return bool(doi) and doi.lower().startswith("10.48550/arxiv")
# ----------------------------------------------------------------- parsers (pure: testable offline)
def parse_arxiv(atom):
ns = {"a": "http://www.w3.org/2005/Atom", "x": "http://arxiv.org/schemas/atom"}
out = []
for e in ET.fromstring(atom).findall("a:entry", ns):
aid = ARXIV_RE.search(e.findtext("a:id", "", ns))
out.append({
"arxiv": aid.group(1) if aid else None,
"title": " ".join(e.findtext("a:title", "", ns).split()),
"authors": [a.findtext("a:name", "", ns) for a in e.findall("a:author", ns)],
"year": (e.findtext("a:published", "", ns) or "")[:4],
"abstract": " ".join(e.findtext("a:summary", "", ns).split()),
"doi": (e.findtext("x:doi", None, ns) or None),
"venue": (e.findtext("x:journal_ref", None, ns) or None), # authors' own note, e.g. "NeurIPS 2017"
"pdf": f"https://arxiv.org/pdf/{aid.group(1)}" if aid else None,
"source": "arxiv",
})
return out
def parse_openalex(js):
out = []
for w in json.loads(js).get("results", []):
ids = w.get("ids", {}) or {}
doi = (w.get("doi") or "").replace("https://doi.org/", "") or None
loc = w.get("primary_location") or {}
src = loc.get("source") or {}
best = w.get("best_oa_location") or {}
# abstract comes as an inverted index
inv = w.get("abstract_inverted_index") or {}
words = sorted((p, t) for t, ps in inv.items() for p in ps)
out.append({
"arxiv": next((m.group(1) for u in (best.get("pdf_url"), best.get("landing_page_url"))
if u and "arxiv.org" in u and (m := ARXIV_RE.search(u))), None),
"title": w.get("title") or "",
"authors": [a.get("author", {}).get("display_name", "") for a in w.get("authorships", [])],
"year": str(w.get("publication_year") or ""),
"abstract": " ".join(t for _, t in words),
"doi": doi,
"venue": src.get("display_name"),
"type": w.get("type"), # article, preprint, book-chapter, ...
"citations": w.get("cited_by_count", 0),
"pdf": best.get("pdf_url"),
"source": "openalex",
})
return out
def parse_s2(js):
data = json.loads(js)
items = data.get("data", [data]) if isinstance(data, dict) else data
out = []
for p in items:
if not p:
continue
ext = p.get("externalIds") or {}
pv = p.get("publicationVenue") or {}
out.append({
"arxiv": ext.get("ArXiv"),
"title": p.get("title") or "",
"authors": [a.get("name", "") for a in p.get("authors", [])],
"year": str(p.get("year") or ""),
"abstract": p.get("abstract") or "",
"doi": ext.get("DOI"),
"venue": p.get("venue") or pv.get("name"),
"venue_type": pv.get("type"), # journal / conference
"types": p.get("publicationTypes") or [], # JournalArticle, Conference, ...
"citations": p.get("citationCount", 0),
"pdf": (p.get("openAccessPdf") or {}).get("url"),
"source": "s2",
})
return out
def parse_crossref(js, want_title):
"""Best Crossref item whose title matches (journal/proceedings only), else None."""
want = norm_title(want_title)
for it in json.loads(js).get("message", {}).get("items", []):
title = (it.get("title") or [""])[0]
if norm_title(title) != want or it.get("type") not in ("journal-article", "proceedings-article", "book-chapter"):
continue
date = (it.get("published-print") or it.get("published-online") or it.get("issued") or {}).get("date-parts", [[None]])[0][0]
return {
"doi": it.get("DOI"), "title": title, "year": str(date or ""),
"venue": (it.get("container-title") or [None])[0] or (it.get("event") or {}).get("name"),
"type": it.get("type"),
"authors": [f"{a.get('given', '')} {a.get('family', '')}".strip() for a in it.get("author", [])],
"source": "crossref",
}
return None
def parse_dblp(js, want_title):
"""dblp hit for the title, preferring the conference/journal record over the CoRR (arXiv) one."""
want = norm_title(want_title)
hits = [h.get("info", {}) for h in json.loads(js).get("result", {}).get("hits", {}).get("hit", [])]
hits = [h for h in hits if norm_title(h.get("title", "").rstrip(".")) == want]
hits.sort(key=lambda h: (h.get("venue") == "CoRR", h.get("type") not in ("Conference and Workshop Papers", "Journal Articles")))
if not hits:
return None
h = hits[0]
return {"key": h.get("key"), "venue": h.get("venue"), "year": h.get("year"), "type": h.get("type"),
"doi": h.get("doi"), "bib": f"https://dblp.org/rec/{h.get('key')}.bib", "source": "dblp"}
def merge(cands):
"""Deduplicate by DOI / arXiv id / normalized title; keep the most informative record."""
seen, out = {}, []
for c in cands:
keys = [k for k in (c.get("doi") and c["doi"].lower(), c.get("arxiv"), norm_title(c["title"])) if k]
hit = next((seen[k] for k in keys if k in seen), None)
if hit is None:
hit = dict(c); out.append(hit)
else:
for k, v in c.items():
if v and not hit.get(k):
hit[k] = v
hit["citations"] = max(hit.get("citations") or 0, c.get("citations") or 0)
for k in keys:
seen[k] = hit
return out
def arxiv_bibtex(rec):
key = bib_key(rec["authors"], rec["year"], rec["title"])
return (f"@misc{{{key},\n title = {{{rec['title']}}},\n author = {{{' and '.join(rec['authors'])}}},\n"
f" year = {{{rec['year']}}},\n eprint = {{{rec['arxiv']}}},\n archivePrefix = {{arXiv}},\n"
f" note = {{arXiv preprint}},\n url = {{https://arxiv.org/abs/{rec['arxiv']}}}\n}}")
def rekey_bibtex(entry, key):
"""Replace the entry key (doi.org / dblp keys are ugly) and tidy whitespace."""
entry = re.sub(r"^\s*@(\w+)\s*\{\s*[^,]*,", lambda m: f"@{m.group(1).lower()}{{{key},", entry.strip(), count=1)
return re.sub(r"\n\s*\n", "\n", entry)
# ----------------------------------------------------------------- network commands
def search(query, max_n=25, since=None):
q = urllib.parse.quote(query)
cands = []
atom = get(f"https://export.arxiv.org/api/query?search_query=all:{q}&max_results={max_n}&sortBy=relevance", "application/atom+xml")
if atom:
cands += parse_arxiv(atom)
oa = get(f"https://api.openalex.org/works?search={q}&per-page={max_n}&mailto={urllib.parse.quote(UA.split('mailto:')[1][:-1])}")
if oa:
cands += parse_openalex(oa)
time.sleep(1)
s2 = get(f"https://api.semanticscholar.org/graph/v1/paper/search?query={q}&limit={max_n}"
"&fields=title,authors,year,abstract,venue,publicationVenue,externalIds,citationCount,openAccessPdf,publicationTypes")
if s2:
cands += parse_s2(s2)
merged = [c for c in merge(cands) if not since or (c.get("year") or "0") >= str(since)]
merged.sort(key=lambda c: -(c.get("citations") or 0))
for c in merged:
c["key"] = bib_key(c.get("authors", []), c.get("year"), c.get("title"))
Path("candidates.json").write_text(json.dumps(merged, indent=2, ensure_ascii=False), encoding="utf-8")
print(f"{len(merged)} candidates (candidates.json). cites year venue key / title")
for c in merged:
venue = (c.get("venue") or ("arXiv" if c.get("arxiv") else "?"))[:26]
print(f"{c.get('citations') or 0:6d} {c.get('year') or ' '} {venue:26s} {c['key']} {c['title'][:70]}")
return merged
def resolve(ident):
"""Find the peer-reviewed version of a paper given an arXiv id, DOI or title."""
arxiv = ARXIV_RE.fullmatch(ident.strip()) or (ARXIV_RE.search(ident) if "arxiv" in ident.lower() else None)
doi = DOI_RE.search(ident)
rec = None
if arxiv or (doi and not is_arxiv_doi(doi.group(1))):
pid = f"arXiv:{arxiv.group(1)}" if arxiv else doi.group(1)
js = get(f"https://api.semanticscholar.org/graph/v1/paper/{pid}?fields=title,authors,year,venue,publicationVenue,externalIds,citationCount,openAccessPdf,publicationTypes")
rec = parse_s2(js)[0] if js and not js.startswith('{"error"') else None
title = rec["title"] if rec else ident.strip()
result = {"query": ident, "title": title, "preprint": None, "published": None}
if rec and rec.get("arxiv"):
result["preprint"] = {"arxiv": rec["arxiv"], "url": f"https://arxiv.org/abs/{rec['arxiv']}", "pdf": f"https://arxiv.org/pdf/{rec['arxiv']}"}
# 1. Semantic Scholar already knows a non-arXiv DOI + venue?
if rec and rec.get("doi") and not is_arxiv_doi(rec["doi"]) and rec.get("venue"):
result["published"] = {"doi": rec["doi"], "venue": rec["venue"], "year": rec["year"], "via": "semanticscholar",
"type": rec.get("venue_type") or ",".join(rec.get("types", []))}
time.sleep(1)
# 2. Crossref by exact title (journal/proceedings only)
if not result["published"]:
cr = get(f"https://api.crossref.org/works?query.bibliographic={urllib.parse.quote(title)}&rows=5&mailto={urllib.parse.quote(UA.split('mailto:')[1][:-1])}")
hit = parse_crossref(cr, title) if cr else None
if hit:
result["published"] = {"doi": hit["doi"], "venue": hit["venue"], "year": hit["year"], "via": "crossref", "type": hit["type"]}
# 3. dblp (CS): venue + a clean .bib
db = get(f"https://dblp.org/search/publ/api?q={urllib.parse.quote(title)}&format=json&h=10")
hit = parse_dblp(db, title) if db else None
if hit and hit["venue"] != "CoRR":
pub = result["published"] or {}
pub.update({k: v for k, v in {"venue": hit["venue"], "year": hit["year"], "doi": pub.get("doi") or hit.get("doi"),
"dblp": hit["key"], "bib": hit["bib"], "type": hit["type"]}.items() if v})
pub["via"] = pub.get("via") or "dblp"
result["published"] = pub
authors = rec["authors"] if rec else []
result["key"] = bib_key(authors, (result["published"] or {}).get("year") or (rec or {}).get("year"), title)
print(json.dumps(result, indent=2, ensure_ascii=False))
if result["published"]:
print(f"\n-> cite the published version: {result['published'].get('venue')} ({result['published'].get('year')}), "
f"DOI {result['published'].get('doi')}; bibtex: python papers.py bibtex {result['published'].get('doi') or result['published'].get('dblp')}")
elif result["preprint"]:
print(f"\n-> no peer-reviewed version found: cite as arXiv preprint {result['preprint']['arxiv']} and say so in research.md")
return result
def bibtex(ident):
ident = ident.strip()
if ident.startswith("conf/") or ident.startswith("journals/"): # dblp key
entry = get(f"https://dblp.org/rec/{ident}.bib", "text/plain")
rec = None
elif (m := DOI_RE.search(ident)) and not is_arxiv_doi(m.group(1)):
entry = get(f"https://doi.org/{m.group(1)}", "application/x-bibtex")
rec = None
else:
aid = ARXIV_RE.search(ident)
if not aid:
print("give a DOI, an arXiv id or a dblp key", file=sys.stderr); return None
atom = get(f"https://export.arxiv.org/api/query?id_list={aid.group(1)}", "application/atom+xml")
rec = parse_arxiv(atom)[0] if atom else None
entry = arxiv_bibtex(rec) if rec else None
if not entry:
print("no BibTeX found", file=sys.stderr); return None
m_auth = re.search(r"author\s*=\s*[{\"](.+?)[}\"]\s*,?\s*\n", entry, re.S | re.I)
m_year = re.search(r"year\s*=\s*[{\"]?(\d{4})", entry, re.I)
m_title = re.search(r"title\s*=\s*[{\"](.+?)[}\"]\s*,?\s*\n", entry, re.S | re.I)
authors = [a.strip() for a in (m_auth.group(1) if m_auth else "").replace("\n", " ").split(" and ")]
authors = [(" ".join(reversed(a.split(", "))) if ", " in a else a) for a in authors]
key = bib_key(authors, m_year.group(1) if m_year else "", re.sub(r"[{}]", "", m_title.group(1)) if m_title else "")
entry = rekey_bibtex(entry, key)
print(entry)
return entry
def download(ident, out="papers"):
Path(out).mkdir(parents=True, exist_ok=True)
aid = ARXIV_RE.search(ident)
url, name = None, None
if ident.startswith("http"):
url, name = ident, re.sub(r"[^\w.-]", "_", ident.rsplit("/", 1)[-1]) or "paper.pdf"
elif aid:
url, name = f"https://arxiv.org/pdf/{aid.group(1)}", f"arxiv_{aid.group(1).replace('.', '_')}.pdf"
elif (m := DOI_RE.search(ident)):
oa = get(f"https://api.openalex.org/works/https://doi.org/{m.group(1)}")
best = (json.loads(oa).get("best_oa_location") or {}) if oa else {}
url, name = best.get("pdf_url"), re.sub(r"[^\w.-]", "_", m.group(1)) + ".pdf"
if not url:
print("no open-access PDF known for this DOI (try the arXiv version, or the publisher page via the library)"); return None
if not url:
print("give an arXiv id, a DOI or a URL", file=sys.stderr); return None
req = urllib.request.Request(url, headers={"User-Agent": UA, "Accept": "application/pdf"})
try:
with urllib.request.urlopen(req, timeout=60) as r:
data = r.read()
except Exception as e: # noqa: BLE001
print(f"download failed: {e}", file=sys.stderr); return None
if not data.startswith(b"%PDF"):
print(f"not a PDF (got {data[:40]!r}) - probably a landing page or a captcha"); return None
path = Path(out) / name
path.write_bytes(data)
print(f"saved {path} ({len(data) // 1024} KB)")
return path
def text(pdf, pages=None):
try:
import pymupdf as fitz
except ImportError:
try:
import fitz
except ImportError:
print("pip install pymupdf", file=sys.stderr); return None
doc = fitz.open(pdf)
n = min(doc.page_count, pages or doc.page_count)
parts = [f"\n\n===== page {i + 1} =====\n" + doc[i].get_text() for i in range(n)]
out = Path(pdf).with_suffix(".txt")
out.write_text("".join(parts), encoding="utf-8")
print(f"{out}: {n}/{doc.page_count} pages, {sum(len(p) for p in parts)} chars")
return out
def main(argv):
if not argv or argv[0] in ("-h", "--help"):
print(__doc__); return 0
cmd, args = argv[0], argv[1:]
opt = {a.split("=")[0]: True for a in args if a.startswith("--")}
def val(flag, default=None):
if flag in args:
i = args.index(flag); return args[i + 1]
return default
pos = [a for i, a in enumerate(args) if not a.startswith("--") and (i == 0 or not args[i - 1].startswith("--"))]
if cmd == "search":
search(" ".join(pos), int(val("--max", 25)), val("--since"))
elif cmd == "resolve":
resolve(" ".join(pos))
elif cmd == "bibtex":
bibtex(" ".join(pos))
elif cmd == "download":
download(pos[0], val("--out", "papers"))
elif cmd == "text":
text(pos[0], int(val("--pages")) if val("--pages") else None)
else:
print(__doc__); return 2
return 0
if __name__ == "__main__":
sys.exit(main(sys.argv[1:]))