fix-bibtex.py

#!/usr/local/bin/python3

import argparse
import json

import bibtexparser
import requests
import sys
import os

import time
from bibtexparser.bwriter import BibTexWriter

CACHE_FILE_PATH = "/tmp/fix-bibtex.cache"


def doi2bib(doi, cache):
    """
    Return a bibTeX string of metadata for a given DOI.
    Reference: https://gist.github.com/jrsmith3/5513926
    """

    if doi in cache:
        return cache[doi]

    url = "http://dx.doi.org/" + doi

    headers = {"accept": "application/x-bibtex"}
    r = requests.get(url, headers=headers)

    cache[doi] = r.text
    return r.text


MAX_RETRY = 10


def safe_doi2bib(doi, cache):
    for i in range(MAX_RETRY):
        result = doi2bib(doi, cache)
        if result != '' and result is not None:
            return result
        time.sleep(1)

    print("timeout")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Fix BibTex File generated by ReadCube")
    parser.add_argument('bibtex_file', metavar='FILE', type=str, help='BibTex file generated by ReadCube for fixing')
    parser.add_argument('--out', dest='output_bibtex_file', metavar='OUTPUT', type=str, help='output file location',
                        required=True)
    args = parser.parse_args()
    if os.path.exists(CACHE_FILE_PATH):
        with open(CACHE_FILE_PATH) as cache_file:
            cache = json.load(cache_file)
    else:
        cache = dict()

    with open(args.bibtex_file) as bibtex_file:
        bib_database = bibtexparser.load(bibtex_file)
    print('%d BibTex entries loaded!' % len(bib_database.entries), file=sys.stderr)
    i = 1
    j = 1
    for bib_entry in bib_database.entries:
        if 'doi' in bib_entry:
            new_bibtex = safe_doi2bib(bib_entry['doi'].replace("\\", ""), cache)
            new_bib_entry = bibtexparser.loads(new_bibtex).entries[0]
            new_bib_entry['ID'] = bib_entry['ID']
            bib_database.entries_dict[new_bib_entry['ID']] = new_bib_entry
            print('%d / %d BibTex entries fixed!' % (i, j), file=sys.stderr)
            i += 1
        j += 1

    bib_database.entries = [value for key, value in bib_database.entries_dict.items()]

    with open(args.output_bibtex_file, 'w') as bibtex_file:
        bibtex_file.write(BibTexWriter().write(bib_database))

    with open(CACHE_FILE_PATH, 'w') as cache_file:
        json.dump(cache, cache_file)