fix SBT v1 loading (#1468)

sourmash-bio · Apr 16, 2021 · d45baba · d45baba
1 parent 9b94bc5
commit d45baba
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/sourmash/sbt.py b/src/sourmash/sbt.py
@@ -825,6 +825,7 @@ def _load_v1(jnodes, leaf_loader, dirname, storage, *, print_version_warning=Tru
             raise ValueError("Empty tree!")
 
         sbt_nodes = {}
+        sbt_leaves = {}
 
         sample_bf = os.path.join(dirname, jnodes[0]['filename'])
         ksize, tablesize, ntables = extract_nodegraph_info(sample_bf)[:3]
@@ -839,13 +840,14 @@ def _load_v1(jnodes, leaf_loader, dirname, storage, *, print_version_warning=Tru
             if 'internal' in jnode['name']:
                 jnode['factory'] = factory
                 sbt_node = Node.load(jnode, storage)
+                sbt_nodes[i] = sbt_node
             else:
                 sbt_node = leaf_loader(jnode, storage)
-
-            sbt_nodes[i] = sbt_node
+                sbt_leaves[i] = sbt_node
 
         tree = SBT(factory, cache_size=cache_size)
         tree._nodes = sbt_nodes
+        tree._leaves = sbt_leaves
 
         return tree