interpro_tsv2function.py

#!/usr/bin/env python3
"""Parse InterProScan tsv output and report predicted function for each gene/protein.
If fasta file is given, it will also report annotated fasta file (.with_functions.fasta).
cat tsv | python tsv2function.py [fasta]
"""

import sys, argparse
from Bio import SeqIO

parser = argparse.ArgumentParser(description="Parse InterProScan tsv to get functions")
parser.add_argument("-i", "--inFile", dest="inFile", required=True, help="InterProScan tsv file")
args = parser.parse_args()

inFile = args.inFile

protid2function = {}
for l in open(inFile):
    l = l.strip()
    ldata  = l.split('\t')
    if len(ldata)>=12:
        protid, function = ldata[0], ldata[12]
        if function.startswith("NULL"):
            pass
        else:
            if function != '-':
                #add prot to dict
                if protid not in protid2function:
                    protid2function[protid] = set()
                #add function
                protid2function[protid].add(function)

#decide if split by pipe
pipeSplit = False
pipes_in_id = len(list(filter(lambda protid: "|" in protid, protid2function.keys())))
#split protid by pipe if not all protids contain pipe
if len(protid2function)>pipes_in_id:
	pipeSplit = True

outfn = inFile.split('.')[0]+".with_functions.txt"
outfile = open(outfn,'w')
for protid, functions in sorted(protid2function.items()):
# 	if pipeSplit:
# 		for p in protid.split('|'):
# 			print (p, "; ".join(functions),sep='\t',file=outfile)
# 	else:
	print (protid, "; ".join(list(functions)),sep='\t',file=outfile)
outfile.close()
	
sys.stderr.write("#%s unique proteins annotated with %s functions\n" % (len(protid2function),sum(len(x) for x in protid2function.values())))

print ("Saved annotated FastA as: %s"%outfn)