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

@@ -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()