From 0ec4cc69967980ba716cbdd212e8179bb801147f Mon Sep 17 00:00:00 2001
From: Victor Lin <13424970+victorlin@users.noreply.github.com>
Date: Mon, 24 Jul 2023 14:06:25 -0700
Subject: [PATCH] Allow metadata index to be exported

Previously, the metadata index was not available for export as a node
attribute. Copy it as a DataFrame column to fix that.
---
 augur/export_v2.py                                     | 10 ++++++++--
 .../export_v2/cram/metadata-with-accession.t           |  3 +--
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/augur/export_v2.py b/augur/export_v2.py
index 1a7c884ad..56c11de3b 100644
--- a/augur/export_v2.py
+++ b/augur/export_v2.py
@@ -1074,10 +1074,16 @@ def run(args):
 
     if args.metadata is not None:
         try:
-            metadata_file = read_metadata(
+            metadata_df = read_metadata(
                 args.metadata,
                 delimiters=args.metadata_delimiters,
-                id_columns=args.metadata_id_columns).to_dict(orient="index")
+                id_columns=args.metadata_id_columns)
+
+            # Add the index as a column.
+            metadata_df[metadata_df.index.name] = metadata_df.index
+
+            metadata_file = metadata_df.to_dict(orient="index")
+
             for strain in metadata_file.keys():
                 if "strain" not in metadata_file[strain]:
                     metadata_file[strain]["strain"] = strain
diff --git a/tests/functional/export_v2/cram/metadata-with-accession.t b/tests/functional/export_v2/cram/metadata-with-accession.t
index f88367a30..218ce73ac 100644
--- a/tests/functional/export_v2/cram/metadata-with-accession.t
+++ b/tests/functional/export_v2/cram/metadata-with-accession.t
@@ -17,7 +17,6 @@ Run export with metadata that contains "accession".
   {}
 
 Run export with metadata that contains "accession", and use "accession" as the ID column.
-Currently, this results in losing the accession from the node attributes.
 
   $ ${AUGUR} export v2 \
   >  --tree "$TESTDIR/../data/tree.nwk" \
@@ -30,4 +29,4 @@ Currently, this results in losing the accession from the node attributes.
 
   $ python3 "$TESTDIR/../../../../scripts/diff_jsons.py" "$TESTDIR/../data/dataset_with_accession.json" dataset.json \
   >   --exclude-paths "root['meta']['updated']" "root['meta']['maintainers']"
-  {'dictionary_item_removed': [root['tree']['children'][0]['node_attrs']['accession'], root['tree']['children'][1]['children'][0]['node_attrs']['accession'], root['tree']['children'][1]['children'][1]['node_attrs']['accession'], root['tree']['children'][2]['children'][0]['node_attrs']['accession'], root['tree']['children'][2]['children'][1]['node_attrs']['accession'], root['tree']['children'][2]['children'][2]['node_attrs']['accession']]}
+  {}