Generate citations for given PMID(s)

citations.py

#!/usr/bin/env python2
# coding: utf-8
"""\033[1;3;34m
    Generate citations for given PMID(s)
\033[0m"""
import os
import re
import sys
import requests

__author__ = 'suqingdong'
__author_email__ = 'suqingdong@novogene.com'


reload(sys)
sys.setdefaultencoding('utf-8')

HEADERS = {
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/78.0.3904.108 Safari/537.36'
}

def get_cit(pmid, fmt='ama'):

    url = 'https://pubmed.ncbi.nlm.nih.gov/{}/citations/'.format(pmid)
    try:
        res = requests.get(url, headers=HEADERS).json()
        cit = res[fmt]['orig']
        return cit
    except ValueError:
        return None


def main():

    if not args['pmid']:
        pmid_list = sys.stdin.read().split()
    elif os.path.isfile(args['pmid']):
        pmid_list = open(args['pmid']).read().split()
    else:
        pmid_list = re.split(r'[,;]', args['pmid'])


    out = open(args['out'], 'w') if args['out'] else sys.stdout

    with out:
        for pmid in pmid_list:
            cit = get_cit(pmid, fmt=args['format'])
            if cit:
                out.write('{}\t{}\n'.format(pmid, cit))
            else:
                sys.stderr.write('no cit for pmid: {}\n'.format(pmid))


if __name__ == '__main__':

    import argparse
    import textwrap

    epilog = textwrap.dedent('''
        \033[32mexamples:
            %(prog)s 111
            %(prog)s 111 -f nlm
            %(prog)s 111,222,333
            %(prog)s 111,222,333 -o out.txt
            %(prog)s 111,222,333 > out.txt
            %(prog)s pmid.txt
            cat pmid.txt | cut -f 1 | %(prog)s

        \033[34mcontact: {__author__} <{__author_email__}>\033[0m
    ''').format(**locals())

    parser = argparse.ArgumentParser(
        prog='citations',
        description=__doc__,
        epilog=epilog,
        formatter_class=argparse.RawTextHelpFormatter)

    parser.add_argument('pmid', help='the pmid(s) to cit', nargs='?')

    parser.add_argument('-o', '--out', help='the output file [stdout]')
    parser.add_argument(
        '-f',
        '--format',
        help='the format of citations [%(default)s]',
        default='ama',
        choices=['ama', 'apa', 'mla', 'nlm'])

    if (len(sys.argv) == 1) and sys.stdin.isatty():
        parser.print_help()
        exit()

    args = vars(parser.parse_args())

    main()