get_genbank_file4

#!/usr/bin/env python

"""
Description: 
AccNum.txt is a file of accession numbers

Usage: 
GenbankDownload pathToThegenomeFasta file
"""
#********************* MODULES**********************
from __future__ import division
from __future__ import print_function
import sys
import os
import os.path
import sys
import re #regular expressions
import time

# Use the correct module based on the python version.
if sys.version_info[0] == 2:
	import httplib as httplib

else: # More modern python
	import http.client as httplib

# XML parsing used below is common to 2.7 and 3.x
import xml.etree.ElementTree as ET

# Entrez Utils global variable (hostname for NIH EUtils)
eutilsHost = "eutils.ncbi.nlm.nih.gov"


#******************* FUNCTION DEFINITIONS *******************
def get(hostname, uri):
	GET="GET"

	connection = httplib.HTTPSConnection(hostname)

	connection.request(GET, uri)
	response = connection.getresponse()

	data = response.read()

	connection.close()

	return (response.status, response.reason, data)


def EFetch(acc, retmode="text", rettype="fasta", db="nuccore",
		email="ksnp-dev@kissake.net"):

	efetchPrefix = "/entrez/eutils/efetch.fcgi?"
	argSeparator = "&"

	arguments={
		"tool": "kSNP4",
		"email": email,
		"db": db,
		"rettype": rettype,
		"retmode": retmode,
		"id": acc,
	}


	argumentsList = [ "=".join(item) for item in arguments.items()]
	argumentsString = argSeparator.join(argumentsList)
	URI = efetchPrefix +  argumentsString
	# Sleep for 1/3 second; enforce a rate limit of 3 requests per second
	time.sleep(0.34)

	(status, reason, fetched) = get(eutilsHost,URI)

	if status != 200:
		sys.stderr.write("HTTP Error %d: %s\n" % (status, reason))
		sys.exit(1)


	return fetched

def ESearch(searchTerm, db="nuccore", email="ksnp-dev@kissake.net", retmax=100):

	idListParent='IdList'
	idListChildren='Id'

	efetchPrefix = "/entrez/eutils/esearch.fcgi?"
	argSeparator = "&"

	arguments={
		"tool": "kSNP4",
		"email": email,
		"db": db,
		"rettype": 'uilist',
		"retmode": 'xml',
		"term": searchTerm,
		"idtype": "acc",
		"retmax": str(retmax),
	}

	URI = efetchPrefix + argSeparator.join([
		"=".join(item) for item in arguments.items()
	])
	# Sleep for 1/3 of a second here to enforce a rate limit of 3
	time.sleep(0.34)

	(status, reason, searchData) = get(eutilsHost,URI)

	if status != 200:
		sys.stderr.write("HTTP Error %d: %s\n" % (status, reason))
		sys.exit(1)

	searchRoot = ET.fromstring(searchData)
	elementsWithIds = searchRoot.find(idListParent).findall(idListChildren)
	ids = [ element.text for element in elementsWithIds ]

	return ids


#******************* FUNCTION DEFINITIONS *******************
def extractAccNum(line):
	separator = '|'
	prefix = '>'
	genbank = 'gb'
	refseq = 'ref'
	# Fails for text: 'gi|1002344102|gb|LPTZ01001107.1|' output as accnum

	temp = line.split()
	ident = temp[0]
	ident = ident.lstrip(prefix)
	if separator not in ident:
		# This is a basic, accession-number-only identifier
		return(ident)
	else:
		# This is an NCBI standard identifier per
		# https://en.wikipedia.org/wiki/FASTA_format
		# so we assume it has a 'gb' reference.
		idParts = ident.split(separator)
		if genbank in idParts:
			gbAcc = idParts.index(genbank) + 1
			return idParts[gbAcc]
		elif refseq in idParts:
			refAcc = idParts.index(refseq) + 1
			return idParts[refAcc]
		else:
			return None

#********************** Main *************************
#variables
fileExist = ''
pathToFastaGenome = sys.argv[1]
accnum = ''
GenomeFile = ''
seqData = False  # et to True iuf line is sequence data
AccNumList = []

startTime = time.time()


#extract the accession number from the sequence file header
INFILE = open(pathToFastaGenome, 'r')

for line in INFILE:
	if line.startswith('>'):
		AccNumList.append(extractAccNum(line))
		print(AccNumList[-1])
INFILE.close()


OUTFILE = open('genbank_from_NCBI.gbk', 'ab',)

for accnum in AccNumList: #for each accession number
	#find the NCBI ID for the accnum
	EntrezEmail = 'barryghall@gmail.com'
	IDs = ESearch(accnum, db='nuccore', email=EntrezEmail, retmax=100)
	print("Found %d entries"%(len(IDs)))
	if len(IDs) > 0:
		ID = str(IDs[0])
		#print("The ID is ", ID)
		try:
			text = EFetch(db="nuccore", acc=accnum,
					rettype="gb", retmode="text")
			print("Fetched %s"%(accnum))
		except:
			continue
		temp = text.split(b'\n')
		for j in range(len(temp)):
			if temp[j].startswith(b"ORIGIN"):
				OUTFILE.write(temp[j])
				OUTFILE.write(b'\n')
				OUTFILE.write(b'//\n\n')
				break
			else:
				# download the accnum file
				OUTFILE.write(temp[j])
				OUTFILE.write(b'\n')

		OUTFILE.write(b'//\n\n')
	time.sleep(1)

OUTFILE.close()
endTime = time.time()
elapsedTime = endTime-startTime
print("Used ",elapsedTime, "seconds.")