Skills › Scientific Agent Skills › citation-management › search_pubmed.py
search_pubmed.py
434 lines · skills/citation-management/scripts/search_pubmed.py at 5f2bc03
1#!/usr/bin/env python3
2"""
3PubMed Search Tool
4Search PubMed using E-utilities API and export results.
5"""
7import sys
8import os
9import re
10import requests
11import argparse
12import json
13import time
14import xml.etree.ElementTree as ET
15from typing import List, Dict, Optional
16from datetime import datetime
18sys.path.insert(0, str(__import__("pathlib").Path(__file__).resolve().parent))
20from _common import ( # noqa: E402
21 citation_key,
22 format_pages,
23 protect_title,
24 render_entry,
25)
27class PubMedSearcher:
28 """Search PubMed using NCBI E-utilities API."""
30 def __init__(self, api_key: Optional[str] = None, email: Optional[str] = None):
31 """
32 Initialize searcher.
34 Args:
35 api_key: NCBI API key (optional but recommended)
36 email: Email for Entrez (optional but recommended)
37 """
38 self.api_key = api_key or os.getenv('NCBI_API_KEY', '')
39 self.email = email or os.getenv('NCBI_EMAIL', '')
40 self.base_url = 'https://eutils.ncbi.nlm.nih.gov/entrez/eutils/'
41 self.session = requests.Session()
43 # Rate limiting
44 self.delay = 0.11 if self.api_key else 0.34 # 10/sec with key, 3/sec without
46 def search(self, query: str, max_results: int = 100,
47 date_start: Optional[str] = None, date_end: Optional[str] = None,
48 publication_types: Optional[List[str]] = None) -> List[str]:
49 """
50 Search PubMed and return PMIDs.
52 Args:
53 query: Search query
54 max_results: Maximum number of results
55 date_start: Start date (YYYY/MM/DD or YYYY)
56 date_end: End date (YYYY/MM/DD or YYYY)
57 publication_types: List of publication types to filter
59 Returns:
60 List of PMIDs
61 """
62 # Build query with filters
63 if max_results < 1:
64 raise ValueError('max_results must be positive')
65 full_query = f'({query})'
67 # Add date range
68 if date_start or date_end:
69 start = (date_start or '1900').replace('-', '/')
70 end = (date_end or datetime.now().strftime('%Y')).replace('-', '/')
71 full_query += f' AND {start}:{end}[Publication Date]'
73 # Add publication types
74 if publication_types:
75 pub_type_query = ' OR '.join([f'"{pt}"[Publication Type]' for pt in publication_types])
76 full_query += f' AND ({pub_type_query})'
78 print(f'Searching PubMed: {full_query}', file=sys.stderr)
80 # PubMed ESearch exposes only the first 10000 matches in total,
81 # so say so rather than letting the caller believe they got everything.
82 if max_results > 10000:
83 print(
84 f'Note: PubMed ESearch exposes only the first 10000 matches; '
85 f'requesting 10000 rather than {max_results}. Narrow the query '
86 f'or split it by date range for a larger set.',
87 file=sys.stderr,
88 )
89 max_results = 10000
91 # ESearch to get PMIDs
92 esearch_url = self.base_url + 'esearch.fcgi'
93 params = {
94 'db': 'pubmed',
95 'term': full_query,
96 'retmax': max_results,
97 'retmode': 'json',
98 'tool': 'citation-management',
99 }
101 if self.email:
102 params['email'] = self.email
103 if self.api_key:
104 params['api_key'] = self.api_key
106 try:
107 response = self.session.get(esearch_url, params=params, timeout=30)
108 response.raise_for_status()
110 data = response.json()
111 if data.get('error') or data.get('esearchresult', {}).get('errorlist'):
112 raise ValueError('ESearch returned a query error')
113 pmids = data['esearchresult']['idlist']
114 count = int(data['esearchresult']['count'])
116 print(f'Found {count} results, retrieving {len(pmids)}', file=sys.stderr)
118 return pmids
120 except Exception as e:
121 print(f'Error searching PubMed ({type(e).__name__}); request details omitted', file=sys.stderr)
122 return []
124 def fetch_metadata(self, pmids: List[str]) -> List[Dict]:
125 """
126 Fetch metadata for PMIDs.
128 Args:
129 pmids: List of PubMed IDs
131 Returns:
132 List of metadata dictionaries
133 """
134 if not pmids:
135 return []
137 metadata_list = []
139 # Fetch in batches of 200
140 batch_size = 200
141 for i in range(0, len(pmids), batch_size):
142 batch = pmids[i:i+batch_size]
143 print(f'Fetching metadata for PMIDs {i+1}-{min(i+batch_size, len(pmids))}...', file=sys.stderr)
145 efetch_url = self.base_url + 'efetch.fcgi'
146 params = {
147 'db': 'pubmed',
148 'id': ','.join(batch),
149 'retmode': 'xml',
150 'tool': 'citation-management',
151 }
153 if self.email:
154 params['email'] = self.email
155 if self.api_key:
156 params['api_key'] = self.api_key
158 try:
159 time.sleep(self.delay)
160 response = self.session.get(efetch_url, params=params, timeout=60)
161 response.raise_for_status()
163 # Parse XML
164 root = ET.fromstring(response.content)
165 if root.find('.//ERROR') is not None:
166 raise ValueError('EFetch returned an error')
167 articles = root.findall('.//PubmedArticle')
168 batch_metadata = []
170 for article in articles:
171 metadata = self._extract_metadata_from_xml(article)
172 if metadata:
173 batch_metadata.append(metadata)
174 if {record['pmid'] for record in batch_metadata} != set(batch):
175 print('Error: EFetch omitted requested records or returned unsupported book records', file=sys.stderr)
176 raise ValueError('Incomplete EFetch batch')
177 metadata_list.extend(batch_metadata)
179 except Exception as e:
180 print(f'Error fetching metadata for batch ({type(e).__name__}); request details omitted', file=sys.stderr)
181 # Never export a failed batch as a complete bibliography.
182 return []
184 return metadata_list
186 @staticmethod
187 def _element_text(element: Optional[ET.Element]) -> str:
188 """Flatten an element's text, including any inline markup children.
190 PubMed terminates titles with a full stop that is not part of the
191 title; a trailing `?` or `!` is, so only the period is dropped.
192 """
193 if element is None:
194 return ''
195 text = ' '.join(''.join(element.itertext()).split())
196 return text[:-1] if text.endswith('.') else text
198 def _extract_metadata_from_xml(self, article: ET.Element) -> Optional[Dict]:
199 """Extract metadata from PubmedArticle XML element."""
200 try:
201 medline_citation = article.find('.//MedlineCitation')
202 article_elem = medline_citation.find('.//Article')
203 journal = article_elem.find('.//Journal')
205 # Get PMID
206 pmid = medline_citation.findtext('.//PMID', '')
208 # Get DOI
209 doi = None
210 article_ids = article.findall('.//ArticleId')
211 for article_id in article_ids:
212 if article_id.get('IdType') == 'doi':
213 doi = article_id.text
214 break
216 # Get authors
217 authors = []
218 author_list = article_elem.find('.//AuthorList')
219 if author_list is not None:
220 for author in author_list.findall('.//Author'):
221 last_name = author.findtext('.//LastName', '')
222 fore_name = author.findtext('.//ForeName', '')
223 if last_name:
224 if fore_name:
225 authors.append(f'{last_name}, {fore_name}')
226 else:
227 authors.append(last_name)
229 # Get year
230 year = article_elem.findtext('.//Journal/JournalIssue/PubDate/Year', '')
231 if not year:
232 medline_date = article_elem.findtext('.//Journal/JournalIssue/PubDate/MedlineDate', '')
233 if medline_date:
234 year_match = re.search(r'\d{4}', medline_date)
235 if year_match:
236 year = year_match.group()
238 # Structured abstracts split into several labelled AbstractText
239 # elements; taking only the first drops most of the abstract.
240 sections = []
241 for chunk in article_elem.findall('.//Abstract/AbstractText'):
242 text = ''.join(chunk.itertext()).strip()
243 if not text:
244 continue
245 label = chunk.get('Label')
246 sections.append(f'{label}: {text}' if label else text)
247 abstract = ' '.join(sections)
250 metadata = {
251 'pmid': pmid,
252 'doi': doi,
253 # itertext, not findtext: titles carry inline markup such as
254 # <i> and <sup>, and findtext returns only the leading run.
255 'title': self._element_text(article_elem.find('.//ArticleTitle')),
256 'authors': ' and '.join(authors),
257 'journal': journal.findtext('.//Title', ''),
258 'year': year,
259 'volume': journal.findtext('.//JournalIssue/Volume', ''),
260 'issue': journal.findtext('.//JournalIssue/Issue', ''),
261 'pages': article_elem.findtext('.//Pagination/MedlinePgn', ''),
262 'abstract': abstract
263 }
265 return metadata
267 except Exception as e:
268 print(f'Error extracting metadata: {e}', file=sys.stderr)
269 return None
271 def metadata_to_bibtex(self, metadata: Dict) -> str:
272 """Convert metadata to BibTeX format.
274 Uses the shared key scheme rather than appending the PMID, so an entry
275 found here and the same paper found via Crossref or OpenAlex collide
276 and can be deduplicated.
277 """
278 key = citation_key(
279 metadata.get('authors', ''),
280 metadata.get('year', ''),
281 metadata.get('title', ''),
282 )
284 fields = {
285 'author': metadata.get('authors', ''),
286 'title': protect_title(metadata.get('title', '')),
287 'journal': metadata.get('journal', ''),
288 'year': metadata.get('year', ''),
289 'volume': metadata.get('volume', ''),
290 'number': metadata.get('issue', ''),
291 'pages': format_pages(metadata.get('pages')),
292 'doi': metadata.get('doi') or '',
293 }
295 if metadata.get('pmid'):
296 fields['note'] = f'PMID: {metadata["pmid"]}'
298 return render_entry('article', key, fields)
301def main():
302 """Command-line interface."""
303 parser = argparse.ArgumentParser(
304 description='Search PubMed using E-utilities API',
305 epilog='Example: python search_pubmed.py "CRISPR gene editing" --limit 100'
306 )
308 parser.add_argument(
309 'query',
310 nargs='?',
311 help='Search query (PubMed syntax)'
312 )
314 parser.add_argument(
315 '--query',
316 dest='query_arg',
317 help='Search query (alternative to positional argument)'
318 )
320 parser.add_argument(
321 '--query-file',
322 help='File containing search query'
323 )
325 parser.add_argument(
326 '--limit',
327 type=int,
328 default=100,
329 help='Maximum number of results (default: 100)'
330 )
332 parser.add_argument(
333 '--date-start',
334 help='Start date (YYYY/MM/DD or YYYY)'
335 )
337 parser.add_argument(
338 '--date-end',
339 help='End date (YYYY/MM/DD or YYYY)'
340 )
342 parser.add_argument(
343 '--publication-types',
344 help='Comma-separated publication types (e.g., "Review,Clinical Trial")'
345 )
347 parser.add_argument(
348 '-o', '--output',
349 help='Output file (default: stdout)'
350 )
352 parser.add_argument(
353 '--format',
354 choices=['json', 'bibtex'],
355 default='json',
356 help='Output format (default: json)'
357 )
359 parser.add_argument(
360 '--api-key',
361 help='NCBI API key (or set NCBI_API_KEY env var)'
362 )
364 parser.add_argument(
365 '--email',
366 help='Email for Entrez (or set NCBI_EMAIL env var)'
367 )
369 args = parser.parse_args()
371 # Get query
372 query = args.query or args.query_arg
374 if args.query_file:
375 try:
376 with open(args.query_file, 'r', encoding='utf-8') as f:
377 query = f.read().strip()
378 except Exception as e:
379 print(f'Error reading query file: {e}', file=sys.stderr)
380 sys.exit(1)
382 if not query:
383 parser.print_help()
384 sys.exit(1)
385 if args.limit < 1:
386 parser.error('--limit must be positive')
388 # Parse publication types
389 pub_types = None
390 if args.publication_types:
391 pub_types = [pt.strip() for pt in args.publication_types.split(',')]
393 # Search PubMed
394 searcher = PubMedSearcher(api_key=args.api_key, email=args.email)
395 pmids = searcher.search(
396 query,
397 max_results=args.limit,
398 date_start=args.date_start,
399 date_end=args.date_end,
400 publication_types=pub_types
401 )
403 if not pmids:
404 print('No results found', file=sys.stderr)
405 sys.exit(1)
407 # Fetch metadata
408 metadata_list = searcher.fetch_metadata(pmids)
409 if not metadata_list:
410 print('Error: incomplete PubMed metadata retrieval; no output written', file=sys.stderr)
411 sys.exit(2)
413 # Format output
414 if args.format == 'json':
415 output = json.dumps({
416 'query': query,
417 'count': len(metadata_list),
418 'results': metadata_list
419 }, indent=2)
420 else: # bibtex
421 bibtex_entries = [searcher.metadata_to_bibtex(m) for m in metadata_list]
422 output = '\n\n'.join(bibtex_entries) + '\n'
424 # Write output
425 if args.output:
426 with open(args.output, 'w', encoding='utf-8') as f:
427 f.write(output)
428 print(f'Wrote {len(metadata_list)} results to {args.output}', file=sys.stderr)
429 else:
430 print(output)
433if __name__ == '__main__':
434 main()