diff --git a/actions/ost-compare-ligand-structures b/actions/ost-compare-ligand-structures index 255ad34b0bc7c474cc8b946238ebdf9757f4f167..7321b29b3a05a8a254222f7bf300a5becc986993 100644 --- a/actions/ost-compare-ligand-structures +++ b/actions/ost-compare-ligand-structures @@ -110,6 +110,7 @@ If BiSyRMSD was enabled with --rmsd, the following columns are added: import argparse import csv +from io import StringIO import json import os import sys @@ -212,6 +213,22 @@ def _ParseArgs(): "reports one model ligand, instead of a reference ligand. " "Has no effect with JSON output.")) + parser.add_argument( + "--csv-extra-header", + dest="csv_extra_header", + default=None, + type=str, + help=("Extra header prefix for CSV output. This allows adding " + "additional annotations (such as target ID, group, etc) to the " + "output")) + + parser.add_argument( + "--csv-extra-data", + dest="csv_extra_data", + default=None, + type=str, + help=("Additional data (columns) for CSV output.")) + parser.add_argument( "-mb", "--model-biounit", @@ -796,6 +813,18 @@ def _WriteCSV(out, args): "rmsd_unassigned": reason[0], }) + if args.csv_extra_header or args.csv_extra_data: + + extra_csv = StringIO( + args.csv_extra_header + os.linesep + args.csv_extra_data) + reader = csv.DictReader(extra_csv) + extra_data = next(iter(reader)) + if None in extra_data: + raise ValueError("Not enough columns in --csv-extra-header") + fieldnames = reader.fieldnames + fieldnames + for ligand, row in csv_dict.items(): + row.update(extra_data) + with open(args.output, 'w', newline='') as csvfile: writer = csv.DictWriter(csvfile, fieldnames=fieldnames) writer.writeheader()