#!/usr/bin/env python3 import json import sys import os import pandas as pd from tabulate import tabulate from argparse import ArgumentParser # JSON class class WriteResultFile: """Recreation of all ResFinder result files from ResFinder json file.""" def __init__(self, json_file, db_path, outdir, sample, mp): self.json_file = json_file self.db_point_path = db_path self.outdir = outdir self.sample = sample self.mapping_method = mp # Load json data and check format self.data = WriteResultFile.open_and_read_json(self) WriteResultFile.check_json_std_format(self) # Load global variable self.species = self.data.get("provided_species") self.gene_lst= WriteResultFile.point_gene_list(self) def parseJSON(self): """ Iterating through seq_regions and seq_variations classes in json file to extract information for result files. All Resfinder result files will be written in this function. """ # Extract databases to list self.db_list = WriteResultFile.extract_databases(self) # Extract phenotypes to dict self.phenotype_dict,self.pheno_dict,self._pheno_species_dict = WriteResultFile.read_phenotypes(self) # Initiate result files self.ResFinder_results_tab = open(self.outdir + "/ResFinder_results_tab.txt", "w") self.ResFinder_results_tab.write("Resistance gene\tIdentity\tAlignment Length/Gene Length\tCoverage\t" + "Position in reference\tContig or Depth\tPosition in contig\tPhenotype\tAccession no.\tNotes\n") self.ResFinder_disinf_results_tab = open(self.outdir + "/ResFinder_disinf_results_tab.txt", "w") self.ResFinder_disinf_results_tab.write("Resistance gene\tIdentity\tAlignment Length/Gene Length\tCoverage\t" + "Position in reference\tContig or Depth\tPosition in contig\tPhenotype\tAccession no.\tNotes\n") self.PointFinder_results_tab = open(self.outdir + "/PointFinder_results.txt", "w") self.PointFinder_results_tab.write("Mutation\tIdentity\tNucleotide change\tAmino acid change\tResistance\tPMID\n") self.ResFinder_phenotable = open(self.outdir + "/pheno_table.txt", "w") self.ResFinder_phenotable.write("# ResFinder phenotype results.\n") self.ResFinder_phenotable = open(self.outdir + "/pheno_table.txt", "w") self.ResFinder_phenotable.write("# ResFinder phenotype results.\n") self.ResFinder_phenotable_species = open(self.outdir + "/pheno_table_" + self.species +".txt", "w") self.ResFinder_phenotable_species.write("# ResFinder phenotype results for " + self.species + ".\n") # Iterate through the json seq regions class self.res_matrix = [] self.point_gene_hits = set() for j in self.data["seq_regions"].items(): self.data_row = [] if j[1].get("ref_database")[0] == self.db_list[0]: # pointfinder hit self.point_gene_hits.add(j[1].get("name")) elif j[1].get("ref_database")[0] == self.db_list[1] or j[1].get("ref_database")[0] == self.db_list[2]: # Add pheno_table data from Resfinder results if float(j[1].get("identity")) == 100.0: if j[1].get("alignment_length") == j[1].get("ref_seq_lenght"): for phenotype in j[1].get("phenotypes"): if phenotype in self.pheno_species_dict: self.pheno_species_dict[phenotype][2].append(3) self.pheno_species_dict[phenotype][3].append("{} ({})".format(j[1].get("name"),j[1].get("ref_id"))) self.pheno_dict[phenotype][2].append(3) self.pheno_dict[phenotype][3].append("{} ({})".format(j[1].get("name"),j[1].get("ref_id"))) elif j[1].get("alignment_length") < j[1].get("ref_seq_lenght"): for phenotype in j[1].get("phenotypes"): if phenotype in self.pheno_species_dict: self.pheno_species_dict[phenotype][2].append(1) self.pheno_species_dict[phenotype][3].append("{} ({})".format(j[1].get("name"),j[1].get("ref_id"))) self.pheno_dict[phenotype][2].append(1) self.pheno_dict[phenotype][3].append("{} ({})".format(j[1].get("name"),j[1].get("ref_id"))) elif float(j[1].get("identity")) < 100.0: if j[1].get("alignment_length") == j[1].get("ref_seq_lenght"): for phenotype in j[1].get("phenotypes"): if phenotype in self.pheno_species_dict: self.pheno_species_dict[phenotype][2].append(2) self.pheno_species_dict[phenotype][3].append("{} ({})".format(j[1].get("name"),j[1].get("ref_id"))) self.pheno_dict[phenotype][2].append(2) self.pheno_dict[phenotype][3].append("{} ({})".format(j[1].get("name"),j[1].get("ref_id"))) elif j[1].get("alignment_length") < j[1].get("ref_seq_lenght"): for phenotype in j[1].get("phenotypes"): if phenotype in self.pheno_species_dict: self.pheno_species_dict[phenotype][2].append(1) self.pheno_species_dict[phenotype][3].append("{} ({})".format(j[1].get("name"),j[1].get("ref_id"))) self.pheno_dict[phenotype][2].append(1) self.pheno_dict[phenotype][3].append("{} ({})".format(j[1].get("name"),j[1].get("ref_id"))) # Extract informations for tab and txt self.data_row.append(j[1].get("name")) self.data_row.append(j[1].get("identity")) self.data_row.append(str(j[1].get("alignment_length")) + ".."+ str(j[1].get("ref_seq_lenght"))) self.data_row.append(j[1].get("coverage")) self.data_row.append(str(j[1].get("ref_start_pos")) + ".."+ str(j[1].get("ref_end_pos"))) if self.mapping_method == "kma": self.data_row.append(j[1].get("depth")) elif self.mapping_method == "blast": self.data_row.append(j[1].get("query_id")) self.data_row.append(str(j[1].get("query_start_pos")) + ".."+ str(j[1].get("query_end_pos"))) phenotypes = set() for pheno in j[1].get("phenotypes"): if pheno in self.phenotype_dict: phenotypes.add(str(self.phenotype_dict.get(j[1].get("phenotypes")[0]))[2:-2].capitalize()) for i in phenotypes: self.data_row.append(i + " resistance") self.data_row.append(j[1].get("ref_acc")) self.data_row.append(j[1].get("notes")[0]) if j[1].get("ref_database")[0] == self.db_list[1]: # resfinder hit self.res_matrix += [self.data_row] # Write Resfinder tab result file WriteResultFile.write_tab(self, self.ResFinder_results_tab) elif j[1].get("ref_database")[0] == self.db_list[2]: # disinfectant hit # Write to Disinfectant tab result file WriteResultFile.write_tab(self, self.ResFinder_disinf_results_tab ) else: print("Error: unknown reference database has excluded seq region:" + j[0]) self.ResFinder_results_tab.close() # Write pheno_table files WriteResultFile.write_pheno_table(self, self.pheno_dict, self.ResFinder_phenotable) WriteResultFile.write_pheno_table(self, self.pheno_species_dict, self.ResFinder_phenotable_species) # Write Resfinder result txt WriteResultFile.write_txt(self) # Iterate through the json sequence variations class self.point_matrix = [] self.gene_set = {} for j in self.data["seq_variations"].items(): # Add genes with pointfinder hits self.gene_set.add(j[0].split("_",1)[0]) # Store data if any phenotype self.data_row = [] if j[1].get("phenotypes"): self.data_row.append(j[1].get("genes")[0] + j[1].get("seq_var")) # RNA mutations if j[0].split("_",1)[0] in ("16S","23S"): self.data_row.append(j[1].get("ref_codon").upper() + "->"+ j[1].get("var_codon").upper()) self.data_row.append("RNA mutations") else: # DNA mutations self.data_row.append(j[1].get("codon_change").upper()) self.data_row.append(j[1].get("ref_aa").upper() + "->"+ j[1].get("var_aa").upper()) self.data_row.append(j[1].get("phenotype").capitalize()) self.data_row.append(j[1].get("pmids")) self.point_matrix += [self.data_row] # Write PointFinder tab result file WriteResultFile.write_tab(self, self.PointFinder_results_tab) self.PointFinder_results_tab.close() # Write PointFinder_table.txt result file WriteResultFile.write_point_table(self) def open_and_read_json(self): """Open and load data from json file""" self.f = open(self.json_file, 'r') data = json.load(self.f) self.f.close() return data def check_json_std_format(self): """Checks if all classes in the std ResFinder json format is contained in the json infile.""" json_file_keys = set(self.data.keys()) json_std_format_keys = {"type", "databases", "seq_regions", "seq_variations", "phenotypes", "software_name", "software_version", "software_commit", "run_date", "key", "provided_species" , "result_summary"} if not json_file_keys.union(json_std_format_keys): sys.exit("Error: json file does not follow the std format.") def extract_databases(self): """Store reference databases from databases class. Returns a list of database names""" self.db_list = [] for db in self.data.get("databases").keys(): if db.startswith("P"): # PointFinder DB self.db_list.insert(0,db) elif db.startswith("R"): # ResFinder DB self.db_list.insert(1,db) elif db.startswith("D"): # DisinFinder DB self.db_list.insert(2,db) else: return print("Error: Unknown database in json file: " + db) return self.db_list def read_phenotypes(self): """ Store phenotypes predicted in phneotype class in dicts Returns 3 dicts: phenotype_dict for Pointfinder and Resfinder hits pheno_dict for pheno table overview self.pheno_species_dict for species specific pheno table overview """ self.phenotype_dict = {} self.pheno_dict = {} self.pheno_species_dict = {} for j in self.data["phenotypes"].items(): if j[1].get("amr_resistant"): # store data for pheno_table if j[1].get("amr_species_relevant"): self.pheno_species_dict[j[1].get("amr_resistance")] = [j[1].get("amr_classes")[0],"Resistant",[],[]] self.pheno_dict[j[1].get("amr_resistance")] = [j[1].get("amr_classes")[0],"Resistant",[],[]] else: self.pheno_dict[j[1].get("amr_resistance")] = [j[1].get("amr_classes")[0],"Resistant",[],[]] if j[1].get("ref_database") == self.db_list[0]: # store phenotype for pointfinder results for k in range(len(j[1].get("seq_variations"))): self.phenotype_dict[j[1].get("seq_variations")[k]] = j[1].get("amr_resistance") else: # store phenotype self.phenotype_dict[j[1].get("amr_resistance")] = j[1].get("amr_classes") else: # store data for pheno_table if j[1].get("amr_species_relevant"): self.pheno_species_dict[j[1].get("amr_resistance")] = [j[1].get("amr_classes")[0],"No resistance",[0],[]] self.pheno_dict[j[1].get("amr_resistance")] = [j[1].get("amr_classes")[0],"No resistance",[0],[]] else: self.pheno_dict[j[1].get("amr_resistance")] = [j[1].get("amr_classes")[0],"No resistance",[0],[]] return self.phenotype_dict, self.pheno_dict, self.pheno_species_dict def group_matrix(self): """Group Pointfinder hits by genes""" self.res = {i: [] for i in self.gene_lst} for row in self.point_matrix: self.res[row[0].split()[0]].append(row) return self.res def write_point_table(self): self.PointFinder_table = open(self.outdir + "/PointFinder_table.txt", "w") self.PointFinder_table.write("Chromosomal point mutations - Results\nSpecies: {}\nGenes: {}" + "\nMapping methode: {}\n\n\nKnown Mutations\n\n".format(self.species,", ".join(sorted(self.gene_lst)),self.mapping_method)) res = WriteResultFile.group_matrix(self) for gene in res.keys(): self.PointFinder_table.write(gene + "\n") if self.res.get(gene): self.PointFinder_table.write("Mutation\tIdentity\tNucleotide change\tAmino acid change\tResistance\tPMID\n") WriteResultFile.write_tab(self, self.PointFinder_table) self.PointFinder_table.write("\n") elif gene in self.point_gene_hits and gene not in self.gene_set: self.PointFinder_table.write("No known mutations found in {}\n\n".format(gene)) else: self.PointFinder_table.write("No hit found\n\n") self.PointFinder_table.close() def write_pheno_table(self, dict, outfile): """Write pheno table files from dict.""" outfile.write("#\n# " + self.sample + "\n" + "#The phenotype 'No resistance' should be interpreted with\n" + "#caution, as it only means that nothing in the used\n" + "# database indicate resistance, but resistance could exist\n" + "# from 'unknown' or not yet implemented sources.\n#\n" + "# The 'Match' column stores one of the integers 0, 1, 2, 3.\n" + "# 0: No match found\n" + "# 1: Match < 100% ID AND match length < ref length or Match = 100% ID AND match length < ref length\n" + "# 2: Match < 100% ID AND match length = ref length\n" + "# 3: Match = 100% ID AND match length = ref length\n" + "# If several hits causing the same resistance are found,\n" + "# the highest number will be stored in the 'Match' column.\n\n" + "# Antimicrobial\tClass\tWGS-predicted phenotype\tMatch\tGenetic background\n") for phenotype in dict.keys(): values = dict.get(phenotype) outfile.write("{}\t{}\t{}\t{}\t{}\n".format(phenotype,values[0],values[1],max(values[2]),(", ".join(sorted(values[3]))))) outfile.close() def write_tab(self, outfile): """Write a list into a tabseperated file.""" outfile.write("\t".join(map(str,self.data_row)) + "\n") def write_txt(self): """Write ResFinder_results.txt from a matrix A=[[]] grouped by phneotypes.""" df = pd.DataFrame(self.res_matrix, columns=["Resistance gene","Identity", "Alignment Length/Gene Length","Coverage","Position in reference", "Contig or Depth","Position in contig","Phenotype","Accession no.","Notes"]) # group by phenotype grouped_df = df.sort_values(by=['Phenotype']).groupby(['Phenotype']) txt_str = "" for key, item in grouped_df: title = key headers = list(grouped_df.get_group(title).columns) rows = list(grouped_df.get_group(title).values) table = WriteResultFile.text_table(title, headers, rows) txt_str += table # print entire table with open(self.outdir + "/ResFinder_results.txt", "w") as f: f.write(txt_str) def text_table(title, headers, rows, table_format='psql'): ''' Create text table USAGE: >>> from tabulate import tabulate >>> title = 'My Title' >>> headers = ['A','B'] >>> rows = [[1,2],[3,4]] >>> print(text_table(title, headers, rows)) +-----------+ | My Title | +-----+-----+ | A | B | +=====+=====+ | 1 | 2 | | 3 | 4 | +-----+-----+ ''' # Create table table = tabulate(rows, headers, tablefmt=table_format) # Prepare title injection width = len(table.split('\n')[0]) tlen = len(title) if tlen + 4 > width: # Truncate oversized titles tlen = width - 4 title = title[:tlen] spaces = width - 2 - tlen left_spacer = ' ' * int(spaces / 2) right_spacer = ' ' * (spaces - len(left_spacer)) # Update table with title table = '\n'.join(['+%s+' % ('-' * (width - 2)), '|%s%s%s|' % (left_spacer, title, right_spacer), table, '\n']) return table def point_gene_list(self): """Extract genes from resistance overview in Pointfinder for relevant species. A list with genes is returned""" gene_set = set() file = os.path.join(self.db_point_path, self.species, "resistens-overview.txt") with open(file) as f: for line in f: if line.startswith("#"): continue gene = line.split()[0].split("_", maxsplit=1)[0] gene_set.add(gene) self.gene_lst = list(gene_set) return self.gene_lst ########################################################################## # PARSE COMMAND LINE OPTIONS ########################################################################## parser = ArgumentParser() parser.add_argument("-i", "--infile", help="JSON file from ResFinder", required=True) parser.add_argument("-o", "--outdir", help="Path to output directory", default=".") parser.add_argument("-db_point", "--db_path_point", help=("Path to the databases for PointFinder."), required=True, default=None) parser.add_argument("-mp", "--mapping_method", help="Mapping method used for ResFinder run, either KMA or BLAST", required=True) parser.add_argument("-s", "--sample_name", help="Name of sample file ran with ResFinder", required=True) args = parser.parse_args() # Load and check input file(s) if args.infile is None: sys.exit("Input Error: No input file provided!\n") infile = args.infile if not os.path.exists(infile): sys.exit("Input Error: Input file not found at expected location: %s"%(infile)) # Check if valid output directory is provided if args.outdir: outdir = os.path.abspath(args.outdir) if not os.path.exists(outdir): os.makedirs(outdir) else: outdir = os.getcwd() # Check if valid database is provided if args.db_path_point: db_path = os.path.abspath(args.db_path_point) if not os.path.exists(args.db_path_point): sys.exit("Input Error: The specified database directory does not " "exist:" + db_path + "\n") # Check existence of the resistens-overview file in a species res_overview_file = "%s/campylobacter/resistens-overview.txt" % (db_path) if not os.path.isfile(res_overview_file): sys.exit("Resistance-overview.txt not found at expected location: %s"%(res_overview_file)) # Defining varibales method_path = args.mapping_method sample = args.sample_name # Call JSON class to recreate ResFinder result files results = WriteResultFile(infile, db_path, outdir, sample, method_path) results.check_json_std_format() results.parseJSON()