Implement DOI fetching for PMC IDs and refactor citation generation

fetch_arxiv_data.py

@@ -6,7 +6,11 @@ HEADERS = {
 }
 
 def fetch_pmc_doi(pmc_id):
-    pass
+    url = f"https://www.ncbi.nlm.nih.gov/pmc/utils/idconv/v1.0/[email protected]&ids={pmc_id}&format=json"
+    response = requests.get(url, headers=HEADERS).json()
+    if response['status'] == 'ok':
+        doi = response['records'][0]['doi']
+        return f"https://doi.org/{doi}"
 
 def fetch_arxiv_doi(arxiv_id):
     page_url = f"https://arxiv.org/abs/{arxiv_id}"
@@ -15,10 +19,18 @@ def fetch_arxiv_doi(arxiv_id):
     doi = page_data.find('td', {'class': "tablecell arxivdoi"}).find('a', {'id': 'arxiv-doi-link'}).text
     return doi
 
-def generate_citation(doi):
+def fetch_citation(doi):
     citation_content = requests.get(doi, headers={ 'User-Agent':HEADERS['User-Agent'], 'Accept': 'text/x-bibliography; style=apa'}).content
     return citation_content.decode('utf-8')
 
-doi = fetch_arxiv_doi("2412.14338")
-citation = generate_citation(doi)
-print(citation)
+def generate_citation(id):
+    if id.startswith('PMC'):
+        doi = fetch_pmc_doi(id)
+    else:
+        doi = fetch_arxiv_doi(id)
+    citation = fetch_citation(doi).replace('\n', ' ').replace('<i>', '').replace('</i>', '')
+    return citation
+
+if __name__ == '__main__':
+    citation = generate_citation('2412.14338')
+    print(citation)