--- a/tools/align_back_trans/align_back_trans.py	Wed Feb 01 06:55:27 2017 -0500
+++ b/tools/align_back_trans/align_back_trans.py	Wed May 10 12:05:28 2017 -0400
@@ -18,12 +18,14 @@
 """
 
 import sys
-from Bio.Seq import Seq
-from Bio.Alphabet import generic_protein
+
+from Bio import AlignIO
+from Bio import SeqIO
+
 from Bio.Align import MultipleSeqAlignment
-from Bio import SeqIO
-from Bio import AlignIO
+from Bio.Alphabet import generic_protein
 from Bio.Data.CodonTable import ambiguous_generic_by_id
+from Bio.Seq import Seq
 
 if "-v" in sys.argv or "--version" in sys.argv:
     print "v0.0.7"
@@ -130,19 +132,23 @@
 def alignment_back_translate(protein_alignment, nucleotide_records, key_function=None, gap=None, table=0):
     """Thread nucleotide sequences onto a protein alignment."""
     # TODO - Separate arguments for protein and nucleotide gap characters?
-    if key_function is None:
-        key_function = lambda x: x
     if gap is None:
         gap = "-"
 
     aligned = []
-    for protein in protein_alignment:
-        try:
-            nucleotide = nucleotide_records[key_function(protein.id)]
-        except KeyError:
-            raise ValueError("Could not find nucleotide sequence for protein %r"
-                             % protein.id)
-        aligned.append(sequence_back_translate(protein, nucleotide, gap, table))
+    try:
+        if key_function is None:
+            for protein in protein_alignment:
+                nucleotide = nucleotide_records[protein.id]
+                aligned.append(sequence_back_translate(protein, nucleotide, gap, table))
+        else:
+            for protein in protein_alignment:
+                nucleotide = nucleotide_records[key_function(protein.id)]
+                aligned.append(sequence_back_translate(protein, nucleotide, gap, table))
+    except KeyError:
+        raise ValueError("Could not find nucleotide sequence for protein %r"
+                         % protein.id)
+
     return MultipleSeqAlignment(aligned)