/scripts/microbes/harvest_bacteria.py

#!/usr/bin/env python

#Dan Blankenberg



#Harvest Bacteria

#Connects to NCBI's Microbial Genome Projects website and scrapes it for information.

#Downloads and converts annotations for each Genome



import sys, os, time

from urllib2 import urlopen

from urllib import urlretrieve

from ftplib import FTP

from BeautifulSoup import BeautifulSoup

from util import get_bed_from_genbank, get_bed_from_glimmer3, get_bed_from_GeneMarkHMM, get_bed_from_GeneMark



assert sys.version_info[:2] >= ( 2, 4 )



#this defines the types of ftp files we are interested in, and how to process/convert them to a form for our use

desired_ftp_files = {'GeneMark':{'ext':'GeneMark-2.5f','parser':'process_GeneMark'},

                    'GeneMarkHMM':{'ext':'GeneMarkHMM-2.6m','parser':'process_GeneMarkHMM'},

                    'Glimmer3':{'ext':'Glimmer3','parser':'process_Glimmer3'},

                    'fna':{'ext':'fna','parser':'process_FASTA'},

                    'gbk':{'ext':'gbk','parser':'process_Genbank'} }







#number, name, chroms, kingdom, group, genbank, refseq, info_url, ftp_url

def iter_genome_projects( url = "http://www.ncbi.nlm.nih.gov/genomes/lproks.cgi?view=1", info_url_base = "http://www.ncbi.nlm.nih.gov/entrez/query.fcgi?db=genomeprj&cmd=Retrieve&dopt=Overview&list_uids=" ):

    for row in BeautifulSoup( urlopen( url ) ).findAll( name = 'tr', bgcolor = ["#EEFFDD", "#E8E8DD"] ):

        row = str( row ).replace( "\n", "" ).replace( "\r", "" )

        

        fields = row.split( "</td>" )

        

        org_num = fields[0].split( "list_uids=" )[-1].split( "\"" )[0]

        

        name = fields[1].split( "\">" )[-1].split( "<" )[0]

        

        kingdom = "archaea"

        if "<td class=\"bacteria\" align=\"center\">B" in fields[2]:

            kingdom = "bacteria"

            

        group = fields[3].split( ">" )[-1]

        

        info_url = "%s%s" % ( info_url_base, org_num )

        

        org_genbank = fields[7].split( "\">" )[-1].split( "<" )[0].split( "." )[0]

        org_refseq = fields[8].split( "\">" )[-1].split( "<" )[0].split( "." )[0]

        

        #seems some things donot have an ftp url, try and except it here:

        try:

            ftp_url = fields[22].split( "href=\"" )[1].split( "\"" )[0]

        except:

            print "FAILED TO AQUIRE FTP ADDRESS:", org_num, info_url

            ftp_url = None

        

        chroms = get_chroms_by_project_id( org_num )

        

        yield org_num, name, chroms, kingdom, group, org_genbank, org_refseq, info_url, ftp_url



def get_chroms_by_project_id( org_num, base_url = "http://www.ncbi.nlm.nih.gov/entrez/query.fcgi?db=genomeprj&cmd=Retrieve&dopt=Overview&list_uids=" ):

    html_count = 0

    html = None

    while html_count < 500 and html == None:

        html_count += 1

        url = "%s%s" % ( base_url, org_num )

        try:

            html = urlopen( url )

        except:

            print "GENOME PROJECT FAILED:", html_count, "org:", org_num, url

            html = None

            time.sleep( 1 ) #Throttle Connection

    if html is None:

        "GENOME PROJECT COMPLETELY FAILED TO LOAD", "org:", org_num,"http://www.ncbi.nlm.nih.gov/entrez/query.fcgi?db=genomeprj&cmd=Retrieve&dopt=Overview&list_uids="+org_num

        return None

    

    chroms = []

    for chr_row in BeautifulSoup( html ).findAll( "tr", { "class" : "vvv" } ):

        chr_row = str( chr_row ).replace( "\n","" ).replace( "\r", "" )

        fields2 = chr_row.split( "</td>" )

        refseq = fields2[1].split( "</a>" )[0].split( ">" )[-1]

        #genbank = fields2[2].split( "</a>" )[0].split( ">" )[-1]

        chroms.append( refseq )

    

    return chroms



def get_ftp_contents( ftp_url ):

    ftp_count = 0

    ftp_contents = None

    while ftp_count < 500 and ftp_contents == None:

        ftp_count += 1

        try:

            ftp = FTP( ftp_url.split("/")[2] )

            ftp.login()

            ftp.cwd( ftp_url.split( ftp_url.split( "/" )[2] )[-1] )

            ftp_contents = ftp.nlst()

            ftp.close()

        except:

            ftp_contents = None

            time.sleep( 1 ) #Throttle Connection

    return ftp_contents



def scrape_ftp( ftp_contents, org_dir, org_num, refseq, ftp_url ):

    for file_type, items in desired_ftp_files.items():

        ext = items['ext']

        ftp_filename = "%s.%s" % ( refseq, ext )

        target_filename = os.path.join( org_dir, "%s.%s" % ( refseq, ext ) )

        if ftp_filename in ftp_contents:

            url_count = 0

            url = "%s/%s" % ( ftp_url, ftp_filename )

            results = None

            while url_count < 500 and results is None:

                url_count += 1

                try:

                    results = urlretrieve( url, target_filename )

                except:

                    results = None

                    time.sleep(1) #Throttle Connection

            if results is None:

                "URL COMPLETELY FAILED TO LOAD:", url

                return

            

            #do special processing for each file type:

            if items['parser'] is not None:

                parser_results = globals()[items['parser']]( target_filename, org_num, refseq )

        else:

            print "FTP filetype:", file_type, "not found for", org_num, refseq

    #FTP Files have been Loaded

    



def process_FASTA( filename, org_num, refseq ):

    fasta = []

    fasta = [line.strip() for line in open( filename, 'rb' ).readlines()]

    fasta_header = fasta.pop( 0 )[1:]

    fasta_header_split = fasta_header.split( "|" )

    chr_name = fasta_header_split.pop( -1 ).strip()

    accesions = {fasta_header_split[0]:fasta_header_split[1], fasta_header_split[2]:fasta_header_split[3]}

    fasta = "".join( fasta )

    

    #Create Chrom Info File:

    chrom_info_file = open( os.path.join( os.path.split( filename )[0], "%s.info" % refseq ), 'wb+' )

    chrom_info_file.write( "chromosome=%s\nname=%s\nlength=%s\norganism=%s\n" % ( refseq, chr_name, len( fasta ), org_num ) )

    try:

        chrom_info_file.write( "gi=%s\n" % accesions['gi'] )

    except:

        chrom_info_file.write( "gi=None\n" )

    try:

        chrom_info_file.write( "gb=%s\n" % accesions['gb'] )

    except:

        chrom_info_file.write( "gb=None\n" )

    try:

        chrom_info_file.write( "refseq=%s\n" % refseq )

    except:

        chrom_info_file.write( "refseq=None\n" )

    chrom_info_file.close()



def process_Genbank( filename, org_num, refseq ):

    #extracts 'CDS', 'tRNA', 'rRNA' features from genbank file

    features = get_bed_from_genbank( filename, refseq, ['CDS', 'tRNA', 'rRNA'] )

    for feature in features.keys():

        feature_file = open( os.path.join( os.path.split( filename )[0], "%s.%s.bed" % ( refseq, feature ) ), 'wb+' )

        feature_file.write( '\n'.join( features[feature] ) )

        feature_file.close()

    print "Genbank extraction finished for chrom:", refseq, "file:", filename



def process_Glimmer3( filename, org_num, refseq ):

    try:

        glimmer3_bed = get_bed_from_glimmer3( filename, refseq )

    except Exception, e:

        print "Converting Glimmer3 to bed FAILED! For chrom:", refseq, "file:", filename, e

        glimmer3_bed = []

    glimmer3_bed_file = open( os.path.join( os.path.split( filename )[0], "%s.Glimmer3.bed" % refseq ), 'wb+' )

    glimmer3_bed_file.write( '\n'.join( glimmer3_bed ) )

    glimmer3_bed_file.close()



def process_GeneMarkHMM( filename, org_num, refseq ):

    try:

        geneMarkHMM_bed = get_bed_from_GeneMarkHMM( filename, refseq )

    except Exception, e:

        print "Converting GeneMarkHMM to bed FAILED! For chrom:", refseq, "file:", filename, e

        geneMarkHMM_bed = []

    geneMarkHMM_bed_bed_file = open( os.path.join( os.path.split( filename )[0], "%s.GeneMarkHMM.bed" % refseq ), 'wb+' )

    geneMarkHMM_bed_bed_file.write( '\n'.join( geneMarkHMM_bed ) )

    geneMarkHMM_bed_bed_file.close()



def process_GeneMark( filename, org_num, refseq ):

    try:

        geneMark_bed = get_bed_from_GeneMark( filename, refseq )

    except Exception, e:

        print "Converting GeneMark to bed FAILED! For chrom:", refseq, "file:", filename, e

        geneMark_bed = []

    geneMark_bed_bed_file = open( os.path.join( os.path.split( filename )[0], "%s.GeneMark.bed" % refseq ), 'wb+' )

    geneMark_bed_bed_file.write( '\n'.join( geneMark_bed ) )

    geneMark_bed_bed_file.close()







def __main__():

    start_time = time.time()

    base_dir = os.path.join( os.getcwd(), "bacteria" )

    try:

        base_dir = sys.argv[1]

    except:

        print "using default base_dir:", base_dir

    

    try:

        os.mkdir( base_dir )

        print "path '%s' has been created" % base_dir

    except:

        print "path '%s' seems to already exist" % base_dir

    

    for org_num, name, chroms, kingdom, group, org_genbank, org_refseq, info_url, ftp_url in iter_genome_projects():

        if chroms is None:

            continue #No chrom information, we can't really do anything with this organism

        #Create org directory, if exists, assume it is done and complete --> skip it

        try:

            org_dir = os.path.join( base_dir, org_num )

            os.mkdir( org_dir )

        except:

            print "Organism %s already exists on disk, skipping" % org_num

            continue

        

        #get ftp contents

        ftp_contents = get_ftp_contents( ftp_url )

        if ftp_contents is None:

            "FTP COMPLETELY FAILED TO LOAD", "org:", org_num, "ftp:", ftp_url

        else:

            for refseq in chroms:

                ftp_result = scrape_ftp( ftp_contents, org_dir, org_num, refseq, ftp_url )

                #FTP Files have been Loaded

                print "Org:", org_num, "chrom:", refseq, "[", time.time() - start_time, "seconds elapsed. ]"

        

        #Create org info file

        info_file = open( os.path.join( org_dir, "%s.info" % org_num ), 'wb+' )

        info_file.write("genome project id=%s\n" % org_num )

        info_file.write("name=%s\n" % name )

        info_file.write("kingdom=%s\n" % kingdom )

        info_file.write("group=%s\n" % group )

        info_file.write("chromosomes=%s\n" % ",".join( chroms ) )

        info_file.write("info url=%s\n" % info_url )

        info_file.write("ftp url=%s\n" % ftp_url )

        info_file.close()

        

    print "Finished Harvesting", "[", time.time() - start_time, "seconds elapsed. ]"

    print "[", ( time.time() - start_time )/60, "minutes. ]"

    print "[", ( time.time() - start_time )/60/60, "hours. ]"



if __name__ == "__main__": __main__()