add --labels-from to sourmash plot

sourmash-bio · ctb · Apr 29, 2023 · Apr 29, 2023 · Apr 29, 2023 · Apr 29, 2023
commit 1a4a6b628807657578f446b64a6ac303fffc12d2
diff --git a/src/sourmash/cli/plot.py b/src/sourmash/cli/plot.py
@@ -49,6 +49,10 @@ def subparser(subparsers):
         help='write clustered matrix and labels out in CSV format (with column'
         ' headers) to this file'
     )
+    subparser.add_argument(
+        '--labels-from', '--labels-load',
+        help='a CSV file containing label information to use on plot',
+    )
 
 
 def main(args):

diff --git a/src/sourmash/commands.py b/src/sourmash/commands.py
@@ -223,16 +223,17 @@ def compare(args):
         notify(f'saving labels to: {labeloutname}')
         with sourmash_args.FileOutputCSV(labeloutname) as fp:
             w = csv.writer(fp)
-            w.writerow(['md5', 'label', 'name', 'filename', 'signature_file'])
+            w.writerow(['order', 'md5', 'label', 'name', 'filename',
+                        'signature_file'])
 
-            for ss, location in siglist:
+            for n, (ss, location) in enumerate(siglist):
                 md5 = ss.md5sum()
                 sigfile = location
                 label = str(ss)
                 name = ss.name
                 filename = ss.filename
 
-                w.writerow([md5, label, name, filename, sigfile])
+                w.writerow([str(n+1), md5, label, name, filename, sigfile])
 
     # output CSV?
     if args.csv:
@@ -271,10 +272,27 @@ def plot(args):
     # not sure how to change this to use f-strings
     notify('...got {} x {} matrix.', *D.shape)
 
-    if args.labeltext:
-        labelfilename = args.labeltext
-    notify(f'loading labels from {labelfilename}')
-    labeltext = [ x.strip() for x in open(labelfilename) ]
+    if args.labels_from:        # use output from --labels-to => labels
+        if args.labeltext:
+            notify("ERROR: cannot supply both --labeltext and --labels-from")
+            sys.exit(-1)
+
+        labelfilename = args.labels_from
+        notify(f"loading labels from CSV file '{labelfilename}'")
+
+        labeltext = []
+        with sourmash_args.FileInputCSV(labelfilename) as r:
+            for row in r:
+                order, label = row['order'], row['label']
+                labeltext.append((int(order), label))
+        labeltext.sort()
+        labeltext = [ t[1] for t in labeltext ]
+    else:                       # use default or --labeltext text file
+        if args.labeltext:
+            labelfilename = args.labeltext
+        notify(f"loading labels from text file '{labelfilename}'")
+        labeltext = [ x.strip() for x in open(labelfilename) ]
+
     if len(labeltext) != D.shape[0]:
         error('{} labels != matrix size, exiting', len(labeltext))
         sys.exit(-1)

diff --git a/tests/test_sourmash.py b/tests/test_sourmash.py
@@ -920,7 +920,7 @@ def test_plot_override_labeltext(runtmp):
 
         print(runtmp.last_result.out)
 
-        assert 'loading labels from new.labels.txt' in runtmp.last_result.err
+        assert "loading labels from text file 'new.labels.txt'" in runtmp.last_result.err
 
         expected = """\
 0\ta
@@ -947,7 +947,7 @@ def test_plot_override_labeltext_fail(runtmp):
     print(runtmp.last_result.out)
     print(runtmp.last_result.err)
     assert runtmp.last_result.status != 0
-    assert 'loading labels from new.labels.txt' in runtmp.last_result.err
+    assert "loading labels from text file 'new.labels.txt'" in runtmp.last_result.err
     assert '3 labels != matrix size, exiting' in runtmp.last_result.err