codon-filter.py

from Bio import SeqIO
import sys


def codonize(seq):
    """Split a string in codons per 3 characters."""
    return [seq[i:i + 3] for i in range(0, len(seq), 3)]


def check_codon(fasta, codon="TTA"):
    """Print fasta record if it contains TTA"""
    i = 0
    for record in SeqIO.parse(fasta, "fasta"):
        if codon in codonize(str(record.seq)):
            print(record.format("fasta"))
            i += 1
    print("Found {} records with at least one inframe occurence of {}".format(i, codon), file=sys.stderr)


if __name__ == '__main__':
    check_codon(sys.argv[1])