get_from_NCBI.py

import argparse, subprocess, sys, fnmatch
from ftplib import FTP

#initiate parser
parser = argparse.ArgumentParser(description="This program will download bacterial sequencing data from NCBI's FTP server")

#optional arguments
parser.add_argument("-b", "--bacterium", help="name of bacterium (eg. Klebsiella_pneumoniae)", type=str) 
parser.add_argument("-d", "--database", default="refseq", choices = ["refseq", "genbank"], type=str)
parser.add_argument("-t", "--type", default="fasta", help="file type", choices = ["genbank", "fasta", "gff", "feature_table"], type=str)
parser.add_argument("-g", "--genomes", help="Will list all available genomes in specified database", choices=["genbank", "refseq"], type=str)
parser.add_argument("-o", "--outputpath", default=".", help="output directory", type=str)

#initialize FTP server
ftp_site = 'ftp.ncbi.nlm.nih.gov'
ftp = FTP(ftp_site)

#print help if no arguments given
if len(sys.argv) == 1:
	parser.print_help()
	sys.exit()
args = parser.parse_args()

#set default variables
NCBI_database = "refseq"
file_type = "fasta"
	
#change default variable based on command arguments (eg. -t genbank)
if args.database == "refseq":
	NCBI_database = "refseq"
elif args.database == "genbank":
	NCBI_database = "genbank"
	
if args.type == "fasta":
	file_type = "--exclude='*cds_from*' --exclude='*rna_from*' --include='*genomic.fna.gz' --exclude='*'"
elif args.type == "genbank":
	file_type = "--include='*genomic.gbff.gz' --exclude='*'"
elif args.type == "gff":
	file_type = "--include='*genomic.gff.gz' --exclude='*'"
elif args.type == "feature_table":
	file_type = "--include='*feature_table.txt.gz' --exclude='*'"
	
#Have a way to list all available genomes
if args.genomes:
	ftp.login()
	ftp.cwd('genomes/%s/bacteria' % NCBI_database)
	dirs = ftp.nlst()
	for genome in dirs:
		print(genome)
	sys.exit()
	
#Check that bacterium is in NCBI genbank/refseq
if args.bacterium:
	ftp.login()
	ftp.cwd('genomes/%s/bacteria' % NCBI_database)
	dirs = ftp.nlst()
	pattern = args.bacterium
	genome_match = fnmatch.filter(dirs, pattern)
	if len(genome_match) == 0:
		print("Bacterium %s not found in %s -----> Exiting program" % (args.bacterium, NCBI_database))
		sys.exit(1)
else:
	print("")
	print("Please provide bacterium name '-b'")
	print("Exiting program")
	print("")
	sys.exit(1)

#The meat of the program	
try:
	subprocess.call("rsync -Lrtv %s rsync://ftp.ncbi.nlm.nih.gov/genomes/%s/bacteria/%s/latest_assembly_versions/*/ %s" % (file_type, NCBI_database, args.bacterium, args.outputpath), shell=True)
except:
	raise Exception("Failed to download files")