Улучшение кода: изменение заголовков fasta в соответствии с внешним файлом

Question

Мне нужна помощь, чтобы привести в порядок / улучшить мой текущий код. Python не самый сильный, и я действительно мог бы использовать некоторую помощь для улучшения кода. Я оба ищу улучшения в отношении производительности, но также и ради более чистого кода (у меня есть МНОГО предложений if, я думаю, что их можно как-то объединить).

У меня есть fasta-файл (faa) этого формата:

>PLMABFGE_00001 Multidrug resistance ABC transporter ATP-binding and permease protein
MTLIMFIAVPLVVVALLPIMQQSRKIGRKRQDELANFSSDSTSVLSEIRLVKSSNGEPKE
LNVGSNRISSLYGIGVKEAFINSLTQPIINMMMLLFLGILGYGAIRVMNGSLTMGTLVSF
LMNLFQIISPVIIISQFFNELSKTSCSTQRINEMLNEPEEIAQDEQNVDITNKELKFEHV
DFSYKDGKPILHDINVQAKPNTVVAFAGPSGGGKSTIFSLIERFYKPTAGEITIGGENID
EISLEN
>PLMABFGE_00002 Multidrug resistance ABC transporter ATP-binding and permease protein
MPGTIRENLLYGLEREVSDDELWQALDMAYAKNFVQEMDDQLDTQIGERGIKLSGGQRQR
IAIARAFLRNPKILMLDEATASLDSESEAMVQKALSSLMKDRTTLVIAHRLSTIVDADTI
YFIDHGTVSGSGKHEELIKSTPLYAEYVKNQFKK
>PLMABFGE_00003 Putative penicillin-binding protein PbpX
MKHKKQFWLFLTALVIVFCGLFNTQSIKAAYSNKQLRSYAQKTMRKNHLRGSIVIVKNGH
RQTVNIGYGYYKRHLKNGSSKLVYPVGSLQKMVTAAMITQLIYKGKFSQNTKISRWYPLL
KNSSKITVGQLMTHTSGINVSGTESSYGVRFSESGAINWTVARANLSGNTGLNSFNYNNA
NYVLLAGIIRKVTGKSYAANLKSRIIKPLHLCHTYIYKSIPRSKTDAISYLYRNGKNYRS
AAYANANVVSQLPGAGNMFSTAGDYYKIQQGMYNGKILSANQFSYMAHLDSKVNTYSGGF
YLKKGGTLKLAYGNFGDTHFVNWMQLTKDNKNGIVMFLNQPYGSKNHIKSVGYGILKHIK
SGTFIKR
>PLMABFGE_00004 Aspartate-semialdehyde dehydrogenase
MVKEYNVAILGATGAVGRRMIDQLAKSTIPVKSVKLLASSRSAGTVLKFRNQDLTVEETT
PDSFDGIDLVLSSAGGSVSKKFLPEAVKRGAVCVDNTSAFRMEEDVPLVVPEVNPEALKT
HHGIIANPNCSTIQMVMALEPIRRKFGLKQIIVSTYQAASGAGQAAINELKKEAQDYLDG
KDMQKDAQILPTKGDKKHYPLAFNLLPQIDVLEDSGYSHEEWKMIHETKKIMLGDMNAKD
IKVTATCVRVPVPIAHGESVYFTVEDESATAQDIMDAVANFPGVVLQDDIKHQIYPQPIN
AEGKRETFVGRIRPDYENKGAFNMWVVADNLLKGAAWNTVENAEYLVKMNLI

Затем у меня есть «индекс-файл» (gff) этого формата:

##gff-version 3
##sequence-region 001__g__Lactobacillus_no_1 1 62925
##sequence-region 001__g__Lactobacillus_no_2 1 62637
##sequence-region 001__g__Lactobacillus_no_3 1 61639
##sequence-region 001__g__Lactobacillus_no_4 1 61357
001__g__Lactobacillus_no_1  Prodigal:2.6    CDS 206 946 .   +   0   ID=PLMABFGE_00001;eC_number=3.6.3.44;Name=lmrA_1;db_xref=COG:COG1132;gene=lmrA_1;inference=ab initio prediction:Prodigal:2.6,similar to AA sequence:UniProtKB:Q9CHL8;locus_tag=PLMABFGE_00001;product=Multidrug resistance ABC transporter ATP-binding and permease protein
001__g__Lactobacillus_no_1  Prodigal:2.6    CDS 986 1450    .   +   0   ID=PLMABFGE_00002;eC_number=3.6.3.44;Name=lmrA_2;db_xref=COG:COG1132;gene=lmrA_2;inference=ab initio prediction:Prodigal:2.6,similar to AA sequence:UniProtKB:Q9CHL8;locus_tag=PLMABFGE_00002;product=Multidrug resistance ABC transporter ATP-binding and permease protein
001__g__Lactobacillus_no_1  Prodigal:2.6    CDS 1493    2596    .   -   0   ID=PLMABFGE_00003;Name=pbpX_1;db_xref=COG:COG1680;gene=pbpX_1;inference=ab initio prediction:Prodigal:2.6,similar to AA sequence:UniProtKB:O31773;locus_tag=PLMABFGE_00003;product=Putative penicillin-binding protein PbpX
001__g__Lactobacillus_no_1  Prodigal:2.6    CDS 2681    3739    .   -   0   ID=PLMABFGE_00004;eC_number=1.2.1.11;Name=asd_1;db_xref=COG:COG0136;gene=asd_1;inference=ab initio prediction:Prodigal:2.6,similar to AA sequence:UniProtKB:Q04797;locus_tag=PLMABFGE_00004;product=Aspartate-semialdehyde dehydrogenase

Мой код выбирает таксономическое описание (например, 001__g__Lactobacillus_no_1) и добавляет его в соответствующий заголовок после идентификатора заголовка. Кроме того, он вставляет описание в заголовок (TAX = и FUNCTION =) и разделители (|). Итак, для первого заголовка конечный результат выглядит следующим образом:

>PLMABFGE_00001|TAX=001__g__Lactobacillus_no_1|FUNCTION=Multidrug resistance ABC transporter ATP-binding and permease protein

На самом деле у меня есть несколько наборов этих файлов, а реальные файлы имеют длину в миллион строк. Вот пути к каталогам для моих файлов, включая файлы (только для 2 из них):

https://www.dropbox.com/sh/d5vnazkehs2f64e/AABflIS1dL94uCm7nEcFzrK2a?dl=0

Мой код выглядит следующим образом:

#!/usr/bin/env python3
import sys
import fileinput

ID = sys.argv[1]

# Paths
faa_path = r"/genes/"+ID+"/"+ID+".faa"
gff_path = r"/genes/"+ID+"/"+ID+".gff"

tax_dict = dict()

# Put taxonomy and ID information in a dictionary
gff = open(gff_path, "r")
for line in gff:
    if line[0].isdigit(): # Taxonomy names will always start with a digit
        tax = ""
        ID = ""
        index = 0
        ID_flag = False
        tax_flag = True
        for char in line:
            if char.isspace(): # Stop appending to taxonomy when whitespace
                tax_flag = False
            if tax_flag:
                tax+=char
            if line[index:index+3] == "ID=": # Start appending to ID
                ID_flag = True
            if char == ";": # Stop appending to ID
                ID_flag = False
            if ID_flag:
                ID+=char
            index += 1
        tax_dict[ID[3:-1]]=tax

# Modify header of fasta file
for line in fileinput.input(faa_path, inplace = 1):
        if line.startswith(">"): # Look for headers
            ID = ""
            function = ""
            ID_flag = True
            function_flag = False
            final = ""
            for char in line:
                if function_flag:
                    function+=char
                if char.isspace(): # Stop appending to ID and start appending to function
                    ID_flag = False
                    function_flag = True
                if ID_flag:
                    ID+=char
            for key in tax_dict:
                if key == ID[1:-1]: # Find matching IDs
                    final += ID+"|TAX="+tax_dict[key]+"|FUNCTION="+function.strip() # Construct new header
            print(final)
        else:
            print (line.strip())

Спасибо за вашу помощь!