Merge branch adding label_colors arg to Phylo.draw

See GitHub pull request #672 and branch  'phylo-draw-tip-colors'
of https://github.com/bamueh/biopython

Bio/Phylo/_utils.py

@@ -280,7 +280,7 @@ def draw_clade(clade, startcol):
 
 def draw(tree, label_func=str, do_show=True, show_confidence=True,
          # For power users
-         axes=None, branch_labels=None, *args, **kwargs):
+         axes=None, branch_labels=None, label_colors=None, *args, **kwargs):
     """Plot the given tree using matplotlib (or pylab).
 
     The graphic is a rooted tree, drawn with roughly the same algorithm as
@@ -322,6 +322,10 @@ def draw(tree, label_func=str, do_show=True, show_confidence=True,
             But if you would like to alter the formatting of confidence values,
             or label the branches with something other than confidence, then use
             this option.
+        label_colors : dict or callable
+            A function or a dictionary specifying the color of the tip label.
+            If the tip label can't be found in the dict or label_colors is
+            None, the label will be shown in black.
     """
 
     try:
@@ -365,6 +369,20 @@ def format_branch_label(clade):
             "branch_labels must be either a dict or a callable (function)"
         format_branch_label = branch_labels
 
+    # options for displaying label colors.
+    if label_colors:
+        if callable(label_colors):
+            def get_label_color(label):
+                return label_colors(label)
+        else:
+            # label_colors is presumed to be a dict
+            def get_label_color(label):
+                return label_colors.get(label, 'black')
+    else:
+        def get_label_color(label):
+            # if label_colors is not specified, use black
+            return 'black'
+
     # Layout
 
     def get_x_positions(tree):
@@ -446,7 +464,8 @@ def draw_clade(clade, x_start, color, lw):
         label = label_func(clade)
         if label not in (None, clade.__class__.__name__):
             axes.text(x_here, y_here, ' %s' %
-                      label, verticalalignment='center')
+                      label, verticalalignment='center',
+                      color=get_label_color(label))
         # Add label above the branch (optional)
         conf_label = format_branch_label(clade)
         if conf_label:

Tests/test_Phylo_depend.py

@@ -64,6 +64,36 @@ def test_draw(self):
         Phylo.draw(apaf, do_show=False, branch_labels={apaf.root: 'Root'})
         Phylo.draw(apaf, do_show=False, branch_labels=lambda c: c.branch_length)
 
+    def test_draw_with_label_colors_dict(self):
+        """Run the tree layout algorithm with a label_colors argument passed in
+        as a dictionary. Don't display tree."""
+        pyplot.ioff()   # Turn off interactive display
+        dollo = Phylo.read(EX_DOLLO, 'phyloxml')
+        apaf = Phylo.read(EX_APAF, 'phyloxml')
+        label_colors_dollo = {
+            'f_50': 'red',
+            'f_34': 'blue',
+        }
+        label_colors_apaf = {
+            '22_MOUSE': 'red',
+            '18_NEMVE': 'blue',
+        }
+        Phylo.draw(dollo, label_colors=label_colors_dollo, do_show=False)
+        Phylo.draw(apaf, label_colors=label_colors_apaf, do_show=False)
+
+    def test_draw_with_label_colors_callable(self):
+        """Run the tree layout algorithm with a label_colors argument passed in
+        as a callable. Don't display tree."""
+        pyplot.ioff()   # Turn off interactive display
+        dollo = Phylo.read(EX_DOLLO, 'phyloxml')
+        apaf = Phylo.read(EX_APAF, 'phyloxml')
+
+        label_colors_dollo = lambda label: 'r' if label == 'f_50' else 'k'
+        label_colors_apaf = lambda label: 'r'
+
+        Phylo.draw(dollo, label_colors=label_colors_dollo, do_show=False)
+        Phylo.draw(apaf, label_colors=label_colors_apaf, do_show=False)
+
     def test_draw_ascii(self):
         """Tree to Graph conversion."""
         handle = StringIO()