From de5cb3bafe7e2fcbc459b4d53c9195952d98129a Mon Sep 17 00:00:00 2001
From: Xavier Robin <xavier.robin@unibas.ch>
Date: Tue, 27 Aug 2024 10:41:10 +0200
Subject: [PATCH] feat: add extra data to the CSV output
This will allow Andriy's to add target, model and group information
cleanly and directly to the CSV without the need for an extra script.
---
actions/ost-compare-ligand-structures | 29 +++++++++++++++++++++++++++
1 file changed, 29 insertions(+)
diff --git a/actions/ost-compare-ligand-structures b/actions/ost-compare-ligand-structures
index 255ad34b0..7321b29b3 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()
--
GitLab