pubGetPmc

#!/usr/bin/env python

# first load the standard libraries from python
#from sys import *
from __future__ import print_function
import sys, time
import logging, optparse, os, glob, urllib2, tempfile, shutil, csv, re, collections
import subprocess
from os.path import *

# add <scriptDir>/lib/ to package search path
progFile = os.path.abspath(sys.argv[0])
progDir  = os.path.dirname(progFile)
pubToolsLibDir = os.path.join(progDir, "lib")
sys.path.insert(0, pubToolsLibDir)

import pubGeneric, util, pubConf

# === COMMAND LINE INTERFACE, OPTIONS AND HELP ===
parser = optparse.OptionParser("""usage: %prog [options] <outDir> - update current PMC copy from PMC ftp server

Does not check each file on disk for existence, but compares 
the last oa_file_list.txt with the current oa_file_list.txt on the FTP
server.

Creates a file download.log with information what has been added/deleted.

Command to download PMC to some directory using lftp without this script:
(http://lftp.yar.ru/) : 
    lftp -e 'set net:socket-buffer 32000000; connect
    ftp://ftp.ncbi.nlm.nih.gov; mirror --delete-first -c --parallel=8
    --ignore-time /pub/pmc/; quit'

""")

parser.add_option("-d", "--debug", dest="debug", action="store_true", help="show debug messages") 
parser.add_option("-p", "--parallelConnections", dest="parallel", action="store", type="int", help="use lftp for faster, parallel downloads, use X number of connections")
parser.add_option("-c", "--check", dest="check", action="store_true", \
    help="do not use oa_file_list.txt to find the right files, but check file existence on disk")
parser.add_option("", "--justCheck", dest="justCheck", action="store_true", \
    help="just check if all files in the current oa_file_list.txt are really on disk")
parser.add_option("", "--auto", dest="auto", action="store_true", \
    help="automatically set the output directory based on pubConf.extDir")
(options, args) = parser.parse_args()

# ==== FUNCTIONs =====

def parseFileList(path):
    " parse pmc oa_file_list.txt into dict PMCID -> filename "
    if not isfile(path):
        logging.info("Could not find %s, probably first download" % path)
        return {}
    logging.info("Parsing %s" % path)
    data = {}
    ifh = open(path)
    ifh.readline() # skip header
    for line in ifh:
        fields = line.strip('\n').split("\t")
        fname, cit, pmcId = fields[:3]
        data[pmcId] = fname
    return data

def downloadFiles(ftpConn, pmcUrl, theirPmc, missingIds, outDir, connCount):
    """ download pmc files with either a ftp connection (single threadd)
    or with lftp and parallel downloads. If connCount==0, then use single thread.
    """
    # special case for lftp parallel download
    if connCount:
        assert(missingIds!=None) # can't do parallel download with on-disk checking
        fileNames = [theirPmc[pmcId] for pmcId in missingIds]
        pubGeneric.lftpGet(pmcUrl, outDir, fileNames, connCount)
        return

    downloadCount = 0
    alreadyExistCount = 0
    errCount = 0

    downloadIds = missingIds

    if downloadIds==None:
        downloadIds = theirPmc.keys()
        logging.info("scheduling all %d ids for download" % len(downloadIds))

    for pmcId in downloadIds:
        fname = theirPmc[pmcId]
        locPath = join(outDir, fname)
        locDir = dirname(locPath)
        if not isdir(locDir):
            logging.info("Making dir %s" % locDir)
            os.makedirs(locDir)
        if isfile(locPath):
            if missingIds!=None:
                logging.info("File %s already exists" % locPath)
            else:
                logging.debug("File %s already exists" % locPath)
            alreadyExistCount += 1
        else:
            downloadCount += 1
            downloadOk = util.ftpDownload(ftpConn, fname, locPath)
            if not downloadOk:
                logging.error("could not download file %s" % fname)
                errCount +=1
            pubGeneric.appendLog(outDir, "add", pmcId+":"+fname)
    logging.info("%d PMC-IDs: %d already here, downloaded %d, %d skipped due to error" % \
        (len(downloadIds), alreadyExistCount, downloadCount, errCount))

def deleteFiles(ourPmc, staleIds, outDir):
    """ remove files from outDir with pmcId in staleIds, using ourPmc-dict (pmcid->fname) to
    find file names 
    """
    delCount = 0
    for pmcId in staleIds:
        fname = ourPmc[pmcId]
        path = join(outDir, fname)
        logging.debug("Deleting %s" % path)
        try:
            os.remove(path)
            delCount +=1
        except OSError:
            logging.warn("Cannot delete local file %s" % path)
            continue
        pubGeneric.appendLog(outDir, "delete", pmcId+":"+fname)
    logging.info("Deleted %d stale local files" % delCount)

def updatePmc(outDir, checkOnDisk):
    """ compare remote and local oa_file_list.txt and download new files
    ignore the last_updated date, replace oa_file_list.txt once finished

    if checkOnDisk is True, always check files on disk instead of using the oa_file_list.txt comparison
    """

    # download from server to .current.txt
    logging.info("Inferring missing PMC-IDs by comparing oa_file_list.txt")
    newListFname = "oa_file_list.new.txt"
    newListPath = join(outDir, newListFname)
    pmcUrl = "ftp://ftp.ncbi.nlm.nih.gov/pub/pmc"
    ftpConn = util.openFtpConn("ftp.ncbi.nlm.nih.gov", "/pub/pmc", "anonymous", "pubGetPmc@pubtools")
    listName = "oa_file_list.txt"
    logging.info("Downloading %s, file %s to %s" % (pmcUrl, listName, newListPath))
    downloadOk = util.ftpDownload(ftpConn, listName, newListPath)
    if not downloadOk:
        raise Exception("could not download oa_file_list.txt")

    pmcIdxPath = join(outDir, "PMC-ids.csv.gz")
    downloadOk = util.ftpDownload(ftpConn, "PMC-ids.csv.gz", pmcIdxPath)
    if not downloadOk:
        raise Exception("could not download PMC-ID file")

    # compare local txt with current.txt 
    locListPath = join(outDir, "oa_file_list.txt")
    ourPmc = parseFileList(locListPath)
    theirPmc  = parseFileList(newListPath)

    missingIds = set(theirPmc).difference(ourPmc)
    staleIds = set(ourPmc).difference(theirPmc) # stale = on our disk but not on ftp

    logging.debug("Missing here: %d" % len(missingIds))
    logging.debug("Stale here: %d" % len(staleIds))

    if checkOnDisk:
        missingIds = None
        logging.info("Ignoring missing file information, just checking disk")

    downloadFiles(ftpConn, pmcUrl, theirPmc, missingIds, outDir, options.parallel)
    deleteFiles(ourPmc, staleIds, outDir)

    logging.debug("Renaming file list files")
    locListOld = join(outDir, "file_list.old.txt")
    if isfile(locListPath):
        os.rename(locListPath, locListOld)
    os.rename(newListPath, locListPath)

def checkPmc(outDir):
    locListPath = join(outDir, "oa_file_list.txt")
    ourPmc = parseFileList(locListPath)
    missCount = 0
    for pmcId, fileName in ourPmc.iteritems():
        locFname = join(outDir, fileName)
        if not isfile(locFname):
            logging.error("%s does not exist" % locFname)
            missCount +=1
    logging.info("%d files missing" % missCount)

def main(args, options):
    if args==[] and not options.auto:
        parser.print_help()
        exit(1)

    if options.auto:
        outDir = join(pubConf.extDir, "pmc")
    else:
        outDir = args[0]

    if not isdir(outDir):
        print("%s does not exist" % outDir)
        sys.exit(1)
    pubGeneric.setupLogging(progFile, options)
    pubGeneric.setLockFile(outDir, "pubGetPmc")

    if options.justCheck:
        checkPmc(outDir)
    else:
        updatePmc(outDir, options.check)

    pubGeneric.removeLockFiles()

# ----------- MAIN --------------
main(args, options)