filter_utr.py

from Bio import SeqIO

utr_sequences = open(snakemake.input[0], "r")
utr_lines = utr_sequences.readlines()
utr_sequences.close()
trans_utr = dict()
for line in utr_lines:
    if (line.startswith(">")):
        trans_id = line[1:].strip()
        if (trans_id not in trans_utr):
            trans_utr[trans_id] = []
    else:
        trans_utr[trans_id].append(line.strip())

transcripts_filename = snakemake.input[1]
transcripts = SeqIO.index(transcripts_filename, "fasta")

output = []

for transcript in transcripts:
    transcript_id = str(transcripts[transcript].id).split("|")[-1].strip()
    seq = str(transcripts[transcript].seq)
    
    try:
        if transcript_id in trans_utr:
            for utr in trans_utr[transcript_id]:
                pos = seq.find(utr)
                if (pos != -1):
                    seq = seq[:pos] + seq[pos + len(utr):]
            long_seq = seq
        else:
            long_seq_len = 0
            long_seq = ""
            for t in trans_utr:
                s = seq
                seq_len = len(trans_utr[t])
                if (seq_len > long_seq_len):
                    long_seq_len = seq_len
                    for utr in trans_utr[t]:
                        pos = s.find(utr)
                        if (pos != -1):
                            s = s[:pos] + s[pos + len(utr):]
                    long_seq = s
    except:
        long_seq = seq
    
    output.append(">" + str(transcripts[transcript].id))
    output.append(long_seq)

output_file = open(snakemake.output[0],"w+")
output_file.write("\n".join(output))
output_file.close()