add_mtbseq_columns.py

import sys
import argparse
import csv

def main(args):
    with open(args.db) as IN:
        reader = csv.DictReader(IN,delimiter="\t")
        if args.debug:
            sys.stderr.write(str(reader.fieldnames)+"\n")
        for c in args.columns:
            if c not in reader.fieldnames:
                sys.stderr.write(f"\nCan't find the column '{c}' in the database file... Exiting!\n\n")
                quit()
        db = {}
        for row in reader:
            db[(row[args.db_gene].strip(),row[args.db_mutation].strip())] = row

    rows = []
    with open(args.infile) as IN:
        for row in csv.DictReader(IN,delimiter="\t"):
            if row["Subst"]==" ": continue
            gene = row["Gene"]
            mut = row["Subst"].split()[0]
            for c in args.columns:
                if (gene,mut) in db:
                    row[c] = db[(gene,mut)][c]
                else:
                    row[c] = " "
            rows.append(row)

    with open(args.outfile,"w") as OUT:
        writer = csv.DictWriter(OUT,fieldnames=list(rows[0]),delimiter = "\t")
        writer.writeheader()
        writer.writerows(rows)

    
    


parser = argparse.ArgumentParser(description='tbprofiler script',formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('-i','--infile',type=str,help='Input result file',required = True)
parser.add_argument('-d','--db',type=str,help='Database file',required = True)
parser.add_argument('-o','--outfile',type=str,help='Output file name',required = True)
parser.add_argument('-g','--db-gene',type=str,default='Gene ID',help='Column name for the locus tag in the database file')
parser.add_argument('-m','--db-mutation',default='AA change',type=str,help='Column name for the mutation in the database file')
parser.add_argument('-c','--columns',nargs="+",type=str,help='Column to add from the database',required = True)
parser.add_argument('--debug',action='store_true',help='Print debug lines')
parser.set_defaults(func=main)
args = parser.parse_args()
args.func(args)