From bde7c56aef9c35c7efae9dcea78c83224c2b97bc Mon Sep 17 00:00:00 2001
From: Richard Challis <rjchallis@gmail.com>
Date: Tue, 6 Apr 2021 13:43:11 +0100
Subject: [PATCH] Set default traverse limit to class Fixes #64

---
 src/genomehubs/lib/fill.py | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/src/genomehubs/lib/fill.py b/src/genomehubs/lib/fill.py
index e3d77dc5..d042c554 100644
--- a/src/genomehubs/lib/fill.py
+++ b/src/genomehubs/lib/fill.py
@@ -6,7 +6,7 @@
 Usage:
     genomehubs fill [--hub-name STRING] [--hub-path PATH] [--hub-version PATH]
                     [--config-file PATH...] [--config-save PATH]
-                    [--es-host URL...]
+                    [--es-host URL...] [--traverse-limit STRING]
                     [--traverse-infer-ancestors] [--traverse-infer-descendants]
                     [--traverse-infer-both] [--traverse-threads INT]
                     [--traverse-depth INT] [--traverse-root STRING]
@@ -25,6 +25,7 @@
     --traverse-infer-descendants  Flag to enable tree traversal from root to tips.
     --traverse-infer-both         Flag to enable tree traversal from tips to root and
                                   back to tips.
+    --traverse-limit STRING       Maximum rank to ascend to during traversal. [Default: class]
     --traverse-root ID            Root taxon id for tree traversal.
     --traverse-threads INT        Number of threads to use for tree traversal. [Default: 1]
     --traverse-weight STRING      Weighting scheme for setting values during tree
@@ -265,6 +266,7 @@ def set_values_from_descendants(
     taxon_id,
     parent,
     taxon_rank,
+    traverse_limit,
     parents,
     descendant_ranks=None,
     attr_dict=None,
@@ -284,7 +286,7 @@ def set_values_from_descendants(
             traverseable = False
         if not traverseable or taxon_id in limits[key]:
             continue
-        traverse_limit = meta[key].get("traverse_limit", None)
+        traverse_limit = meta[key].get("traverse_limit", traverse_limit)
         if traverse_limit:
             if (
                 descendant_ranks is not None
@@ -333,7 +335,7 @@ def set_values_from_descendants(
     return changed, attr_dict
 
 
-def set_attributes_to_descend(meta):
+def set_attributes_to_descend(meta, traverse_limit):
     """Set which attributes should have values inferred from ancestral taxa."""
     desc_attrs = set()
     desc_attr_limits = {}
@@ -346,6 +348,8 @@ def set_attributes_to_descend(meta):
                 desc_attrs.add(key)
                 if "traverse_limit" in value:
                     desc_attr_limits.update({key: value["traverse_limit"]})
+                else:
+                    desc_attr_limits.update({key: traverse_limit})
     return desc_attrs, desc_attr_limits
 
 
@@ -408,18 +412,15 @@ def traverse_from_tips(es, opts, *, template, root=None, max_depth=None):
         )
     root_depth = max_depth
     meta = template["types"]["attributes"]
-    # for key, value in meta.items():
-    #     if "traverse_limit" in value:
-    #         if not isinstance(value["traverse_limit"], list):
-    #             value["traverse_limit"] = [value["traverse_limit"]]
-    #         value["traverse_limit"] = set(value["traverse_limit"])
     attrs = set(meta.keys())
     parents = defaultdict(
         lambda: defaultdict(lambda: {"max": None, "min": None, "values": []})
     )
     limits = defaultdict(set)
     if "traverse-infer-both" in opts and opts["traverse-infer-both"]:
-        desc_attrs, desc_attr_limits = set_attributes_to_descend(meta)
+        desc_attrs, desc_attr_limits = set_attributes_to_descend(
+            meta, opts["traverse-limit"]
+        )
         missing_attributes = defaultdict(dict)
         descendant_ranks = defaultdict(set)
     else:
@@ -459,6 +460,7 @@ def traverse_from_tips(es, opts, *, template, root=None, max_depth=None):
                     parents=parents,
                     descendant_ranks=descendant_ranks,
                     taxon_rank=node["_source"]["taxon_rank"],
+                    traverse_limit=opts["traverse-limit"],
                     attr_dict=attr_dict,
                     limits=limits,
                 )