Add ability to set mark default.

tanghaibao · Oct 13, 2018 · 5e1f05c · 5e1f05c
1 parent 58db35e
commit 5e1f05c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/goatools/rpt/write_hierarchy_base.py b/goatools/rpt/write_hierarchy_base.py
@@ -10,6 +10,8 @@
 class WrHierPrt(object):
     """Print the hierarchy of a set of objects w/attrs=[item_id, children]."""
 
+    fmt_dashes = '{DASHES} {ID}'
+
     def __init__(self, id2obj, id2nt, cfg, prt=sys.stdout):
         self.id2obj = id2obj  # Contains children (and parents)
         self.id2nt = id2nt    # Contains fields for printing (optional)
@@ -18,6 +20,7 @@ def __init__(self, id2obj, id2nt, cfg, prt=sys.stdout):
         self.max_indent = cfg['max_indent']
         self.include_only = cfg['include_only']
         self.item_marks = self._init_item_marks(cfg.get('item_marks'))
+        self.mark_dflt = cfg.get('mark_default', ' ')
         self.concise_prt = cfg.get('concise_prt', False)
         self.indent = cfg.get('indent', True)
         self.space_branches = cfg.get('space_branches', False)
@@ -43,7 +46,7 @@ def prt_hier_rec(self, item_id, depth=1):
         # Print marks if provided
         if self.item_marks:
             self.prt.write('{MARK} '.format(
-                MARK=self.item_marks[item_id] if item_id in self.item_marks else ' '))
+                MARK=self.item_marks.get(item_id, self.mark_dflt)))
 
         no_repeat = self.concise_prt and item_id in self.items_printed
         # Print content
@@ -69,14 +72,14 @@ def _prtfmt(self, item_id, dashes):
         ntprt = self.id2nt[item_id]
         dct = ntprt._asdict()
         self.prt.write('{DASHES:{N}}'.format(
-            DASHES="{DASHES} {ID}".format(DASHES=dashes, ID=self.nm2prtfmt['ID'].format(**dct)),
+            DASHES=self.fmt_dashes.format(DASHES=dashes, ID=self.nm2prtfmt['ID'].format(**dct)),
             N=self.dash_len))
         self.prt.write("{INFO}\n".format(INFO=self.nm2prtfmt['ITEM'].format(**dct)))
 
     def _prtstr(self, obj, dashes):
         """Print object information using a namedtuple and a format pattern."""
         self.prt.write('{DASHES:{N}}'.format(
-            DASHES="{DASHES} {ID}".format(DASHES=dashes, ID=obj.item_id),
+            DASHES=self.fmt_dashes.format(DASHES=dashes, ID=obj.item_id),
             N=self.dash_len))
         self.prt.write("{INFO}\n".format(INFO=str(obj)))
 

diff --git a/goatools/test_data/nature3102_goea.py b/goatools/test_data/nature3102_goea.py
@@ -53,7 +53,7 @@ def get_goeaobj(method, geneids_pop, taxid):
         propagate_counts=False,
         alpha=0.05,
         methods=[method])
-     # godag is also found in goeaobj.godag
+    # godag is also found in goeaobj.godag
     return goeaobj