add tests for & fix bug in lca_utils.display_lineage

sourmash/lca/lca_utils.py

@@ -12,7 +12,7 @@
 
 __all__ = ['taxlist', 'zip_lineage', 'build_tree', 'find_lca',
            'load_single_database', 'load_databases', 'gather_assignments',
-           'count_lca_for_assignments', 'LineagePair']
+           'count_lca_for_assignments', 'LineagePair', 'display_lineage']
 
 try:                                      # py2/py3 compat
     from itertools import zip_longest
@@ -86,16 +86,25 @@ def zip_lineage(lineage, include_strain=True, truncate_empty=False):
     ['a', '', 'c', '', '', '', '', '']
     """
 
-    row = []
     empty = LineagePair(None, '')
-    for taxrank, lineage_tup in zip_longest(taxlist(include_strain=include_strain), lineage, fillvalue=empty):
-        if lineage_tup == empty:
-            if truncate_empty:
-                break
-        else:
-            # validate non-empty tax, e.g. superkingdom/phylum/class in order.
-            if lineage_tup.rank != taxrank:
-                raise ValueError('incomplete lineage at {} - is {} instead'.format(taxrank, lineage_tup.rank))
+
+    pairs = zip_longest(taxlist(include_strain=include_strain),
+                        lineage, fillvalue=empty)
+    pairs = list(pairs)
+
+    # eliminate empty if so requested
+    if truncate_empty:
+        last_lineage_tup = pairs[-1][1]
+        while pairs and last_lineage_tup == empty:
+            pairs.pop(-1)
+            if pairs:
+                last_lineage_tup = pairs[-1][1]
+
+    row = []
+    for taxrank, lineage_tup in pairs:
+        # validate non-empty tax, e.g. superkingdom/phylum/class in order.
+        if lineage_tup != empty and lineage_tup.rank != taxrank:
+            raise ValueError('incomplete lineage at {} - is {} instead'.format(taxrank, lineage_tup.rank))
 
         row.append(lineage_tup.name)
     return row

tests/test_lca.py

@@ -44,6 +44,11 @@ def test_zip_lineage_3():
     assert zip_lineage(x) == ['a', '', 'c', '', '', '', '', '']
 
 
+def test_zip_lineage_3_truncate():
+    x = [ LineagePair('superkingdom', 'a'), LineagePair(None, ''), LineagePair('class', 'c') ]
+    assert zip_lineage(x, truncate_empty=True) == ['a', '', 'c']
+
+
 def test_zip_lineage_4():
     x = [ LineagePair('superkingdom', 'a'), LineagePair('class', 'c') ]
     with pytest.raises(ValueError) as e:
@@ -52,6 +57,16 @@ def test_zip_lineage_4():
     assert 'incomplete lineage at phylum - is class instead' in str(e.value)
 
 
+def test_display_lineage_1():
+    x = [ LineagePair('superkingdom', 'a'), LineagePair('phylum', 'b') ]
+    assert display_lineage(x) == "a;b", display_lineage(x)
+
+
+def test_display_lineage_2():
+    x = [ LineagePair('superkingdom', 'a'), LineagePair(None, ''), LineagePair('class', 'c') ]
+    assert display_lineage(x) == "a;;c", display_lineage(x)
+
+
 def test_build_tree():
     tree = build_tree([[LineagePair('rank1', 'name1'),
                         LineagePair('rank2', 'name2')]])