#!/usr/bin/env python3 """ ArXiv 论文搜索。通过 ArXiv API（返回 Atom XML）。 支持： - 全文 / 标题 / 摘要 / 作者字段搜索 - 分类过滤、排序 - 按 ID 列表直接拉取论文元数据 - 布尔组合查询（AND / OR / ANDNOT） 示例： python3 arxiv_search.py "attention mechanism" python3 arxiv_search.py "transformer" --category cs.CL --sort date python3 arxiv_search.py "diffusion model" --author "ho jonathan" python3 arxiv_search.py "ViT" --title-only python3 arxiv_search.py --id-list 2409.05591,2301.00001 """ from __future__ import annotations import sys import xml.etree.ElementTree as ET from search_utils import build_parser, get_client, make_item, make_result, print_json API_URL = "https://export.arxiv.org/api/query" # Atom XML 命名空间 NS = { "atom": "http://www.w3.org/2005/Atom", "arxiv": "http://arxiv.org/schemas/atom", } def build_search_query( query: str, category: str | None = None, author: str | None = None, title_only: bool = False, ) -> str: """ 构建 arXiv 查询字符串。 字段前缀： all: 全字段（默认） ti: 仅标题 au: 作者（支持通配 au:smi*） abs: 摘要 cat: 分类 布尔运算符必须大写：AND / OR / ANDNOT """ # 主查询字段 field = "ti" if title_only else "all" parts = [f"{field}:{query}"] if author: # 多个作者用 OR 连接，支持 "lastname firstname" 格式 author_terms = [f"au:{a.strip()}" for a in author.split(",") if a.strip()] if author_terms: parts.append(f"({' OR '.join(author_terms)})") if category: parts.append(f"cat:{category}") return " AND ".join(parts) def fetch_by_ids(id_list: list[str], limit: int) -> list[dict]: """通过 ID 列表直接获取论文元数据（不做文本搜索）。""" params = { "id_list": ",".join(id_list[:limit]), "max_results": min(len(id_list), limit, 100), } with get_client(timeout=30, headers={"Accept": "application/xml"}) as client: resp = client.get(API_URL, params=params) resp.raise_for_status() return _parse_entries(ET.fromstring(resp.text), limit) def search( query: str, limit: int, category: str | None = None, sort_by: str = "relevance", author: str | None = None, title_only: bool = False, ) -> list[dict]: """执行 ArXiv 关键词搜索。""" search_query = build_search_query(query, category, author, title_only) sort_map = { "relevance": "relevance", "date": "lastUpdatedDate", "submitted": "submittedDate", } params = { "search_query": search_query, "start": 0, "max_results": min(limit, 100), "sortBy": sort_map.get(sort_by, "relevance"), "sortOrder": "descending", } with get_client(timeout=30, headers={"Accept": "application/xml"}) as client: resp = client.get(API_URL, params=params) resp.raise_for_status() return _parse_entries(ET.fromstring(resp.text), limit) def _parse_entries(root: ET.Element, limit: int) -> list[dict]: """从 Atom XML 解析论文条目。""" items = [] for entry in root.findall("atom:entry", NS)[:limit]: title = _text(entry, "atom:title").replace("\n", " ").strip() summary = _text(entry, "atom:summary").replace("\n", " ").strip() published = _text(entry, "atom:published") updated = _text(entry, "atom:updated") # 获取论文链接（优先 abs 页面） url = "" pdf_url = "" for link in entry.findall("atom:link", NS): href = link.get("href", "") if link.get("title") == "pdf": pdf_url = href elif link.get("type") == "text/html" or "/abs/" in href: url = href if not url: url = _text(entry, "atom:id") # 从 abs URL 或 id 提取 arxiv_id arxiv_id = "" raw_id = _text(entry, "atom:id") if "/abs/" in raw_id: arxiv_id = raw_id.split("/abs/")[-1] elif raw_id.startswith("http"): arxiv_id = raw_id.split("/")[-1] # 获取作者 authors = [_text(a, "atom:name") for a in entry.findall("atom:author", NS)] # 获取分类 categories = [c.get("term", "") for c in entry.findall("atom:category", NS)] comment = _text(entry, "arxiv:comment") journal_ref = _text(entry, "arxiv:journal_ref") doi = _text(entry, "arxiv:doi") primary_category = entry.find("arxiv:primary_category", NS) primary_cat = primary_category.get("term", "") if primary_category is not None else "" # HTML 版本链接（较新论文有） html_url = f"https://arxiv.org/html/{arxiv_id}" if arxiv_id else None items.append(make_item( title=title, url=url, snippet=summary, arxiv_id=arxiv_id if arxiv_id else None, authors=authors, published=published, updated=updated, pdf_url=pdf_url, html_url=html_url, categories=categories, primary_category=primary_cat if primary_cat else None, comment=comment if comment else None, journal_ref=journal_ref if journal_ref else None, doi=doi if doi else None, )) return items def _text(elem: ET.Element, tag: str) -> str: """安全获取子元素文本。""" child = elem.find(tag, NS) return child.text.strip() if child is not None and child.text else "" def main(): parser = build_parser("搜索 ArXiv 学术论文") parser.add_argument("--category", "-c", help="ArXiv 分类过滤（如 cs.AI, cs.CL, math.CO）") parser.add_argument( "--sort", default="relevance", choices=["relevance", "date", "submitted"], help="排序方式（默认 relevance）", ) parser.add_argument( "--author", "-a", help="按作者过滤（如 'hinton'，多个作者用逗号分隔）", ) parser.add_argument( "--title-only", action="store_true", help="仅在标题中搜索（默认搜索全字段）", ) parser.add_argument( "--id-list", help="直接按 arXiv ID 获取元数据，逗号分隔（如 2409.05591,2301.00001）。指定此项时 query 参数可留空。", ) # 当使用 --id-list 时 query 可选 parser.prog = "arxiv_search.py" # 为了支持 --id-list 时 query 可省略，临时让 query 可选 for action in parser._positionals._group_actions: if action.dest == "query": action.nargs = "?" action.default = "" break args = parser.parse_args() try: if args.id_list: id_list = [i.strip() for i in args.id_list.split(",") if i.strip()] items = fetch_by_ids(id_list, args.limit) query_str = f"id_list:{args.id_list}" else: if not args.query: parser.error("请提供搜索关键词，或使用 --id-list 按 ID 查询") items = search( args.query, args.limit, category=args.category, sort_by=args.sort, author=args.author, title_only=args.title_only, ) query_str = args.query print_json(make_result(True, query_str, "arxiv", items)) except Exception as e: print_json(make_result(False, getattr(args, "query", "") or "", "arxiv", [], str(e))) sys.exit(1) if __name__ == "__main__": main()