#!/usr/bin/env python

###################################################################
##
## gbk2orf.py by Errol Strain (estrain@gmail.com)
##
## Read a GenBank file and export fasta formatted amino acid and 
## CDS files
##
###################################################################

import sys
from optparse import OptionParser
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord


## Command line usage
usage = "usage: %prog -g input.gbk -a aa.fasta -n nuc.fasta" 
p = OptionParser(usage)
p.add_option("-t","--translate", dest="transtabl",type="int",default=11,
  help="Translation table used to translate coding regions (default=11)")
p.add_option("-g","--genbank", dest="gb_file",help="GenBank input file")
p.add_option("-a","--amino_acid", dest="aa_file",help="Fasta amino acid output")
p.add_option("-n","--nucleotide", dest="orf_file",help="Fasta nucleotide output")
(opts, args) = p.parse_args()
## Do I need this next line?
if not opts and not args : p.error("Use --help to see usage")
if len(sys.argv)==1 : p.error("Use --help to see usage") 

## Lists to hold SeqRecords
aalist = []
nuclist = []

## If the CDS does not have a locus tag the name will be assigned using the
## order in which it was found
feat_count=0

## Iterate through genbank records in input file
for gb_record in SeqIO.parse(open(opts.gb_file,"r"), "genbank") :
  for (index, feature) in enumerate(gb_record.features) :
    if feature.type=="CDS" :
      feat_count = feat_count + 1
      gene = feature.extract(gb_record.seq)
      if "locus_tag" in feature.qualifiers :
        value = feature.qualifiers["locus_tag"][0]
      else :
        value =  "Index_" + str(feat_count)
      nuclist.append(SeqRecord(Seq(str(gene)),id=value,name=value))
      pro=Seq(str(gene.translate(table=opts.transtabl,to_stop=True)))
      aalist.append(SeqRecord(pro,id=value,name=value))

## Write out lists in fasta format
aa_handle = open(opts.aa_file,"w")
SeqIO.write(aalist,aa_handle,"fasta")
aa_handle.close()
orf_handle = open(opts.orf_file,"w")
SeqIO.write(nuclist,orf_handle,"fasta")
orf_handle.close()