Merge pull request #15898 from DanHickstein/legend_textcolor

timhoffm · web-flow · commit 8cd780fccd3f · 2020-05-05T23:45:39.000+02:00
New textcolor kwarg for legend
diff --git a/doc/users/next_whats_new/2019-12-09-legend-labelcolor.rst b/doc/users/next_whats_new/2019-12-09-legend-labelcolor.rst
@@ -0,0 +1,16 @@
+Text color for legend labels
+----------------------------
+
+The text color of legend labels can now be set by passing a parameter
+``labelcolor`` to `~.axes.Axes.legend`. The ``labelcolor`` keyword can be:
+
+* A single color (either a string or RGBA tuple), which adjusts the text color
+  of all the labels.
+* A list or tuple, allowing the text color of each label to be set
+  individually.
+* ``linecolor``, which sets the text color of each label to match the
+  corresponding line color.
+* ``markerfacecolor``, which sets the text color of each label to match the
+  corresponding marker face color.
+* ``markeredgecolor``,  which sets the text color of each label to match the
+  corresponding marker edge color.
diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py
@@ -21,13 +21,14 @@
 information.
 """
 
+import itertools
 import logging
 import time
 
 import numpy as np
 
 import matplotlib as mpl
-from matplotlib import cbook, docstring
+from matplotlib import cbook, docstring, colors
 from matplotlib.artist import Artist, allow_rasterization
 from matplotlib.cbook import silent_list
 from matplotlib.font_manager import FontProperties
@@ -172,6 +173,12 @@ def _update_bbox_to_anchor(self, loc_in_canvas):
     absolute font size in points. String values are relative to the current
     default font size. This argument is only used if *prop* is not specified.
 
+labelcolor : str or list
+    Sets the color of the text in the legend. Can be a valid color string
+    (for example, 'red'), or a list of color strings. The labelcolor can
+    also be made to match the color of the line or marker using 'linecolor',
+    'markerfacecolor' (or 'mfc'), or 'markeredgecolor' (or 'mec').
+
 numpoints : int, default: :rc:`legend.numpoints`
     The number of marker points in the legend when creating a legend
     entry for a `.Line2D` (line).
@@ -293,7 +300,8 @@ def __init__(self, parent, handles, labels,
                  scatterpoints=None,    # number of scatter points
                  scatteryoffsets=None,
                  prop=None,          # properties for the legend texts
-                 fontsize=None,        # keyword to set font size directly
+                 fontsize=None,      # keyword to set font size directly
+                 labelcolor=None,    # keyword to set the text color
 
                  # spacing & pad defined as a fraction of the font-size
                  borderpad=None,      # the whitespace inside the legend border
@@ -505,6 +513,36 @@ def __init__(self, parent, handles, labels,
         self.set_title(title, prop=tprop)
         self._draggable = None
 
+        # set the text color
+
+        color_getters = {  # getter function depends on line or patch
+            'linecolor':       ['get_color',           'get_facecolor'],
+            'markerfacecolor': ['get_markerfacecolor', 'get_facecolor'],
+            'mfc':             ['get_markerfacecolor', 'get_facecolor'],
+            'markeredgecolor': ['get_markeredgecolor', 'get_edgecolor'],
+            'mec':             ['get_markeredgecolor', 'get_edgecolor'],
+        }
+        if labelcolor is None:
+            pass
+        elif isinstance(labelcolor, str) and labelcolor in color_getters:
+            getter_names = color_getters[labelcolor]
+            for handle, text in zip(self.legendHandles, self.texts):
+                for getter_name in getter_names:
+                    try:
+                        color = getattr(handle, getter_name)()
+                        text.set_color(color)
+                        break
+                    except AttributeError:
+                        pass
+        elif np.iterable(labelcolor):
+            for text, color in zip(self.texts,
+                                   itertools.cycle(
+                                       colors.to_rgba_array(labelcolor))):
+                text.set_color(color)
+        else:
+            raise ValueError("Invalid argument for labelcolor : %s" %
+                             str(labelcolor))
+
     def _set_artist_props(self, a):
         """
         Set the boilerplate props for artists added to axes.
diff --git a/lib/matplotlib/tests/test_legend.py b/lib/matplotlib/tests/test_legend.py
@@ -534,6 +534,66 @@ def test_legend_title_fontsize():
     assert leg.get_title().get_fontsize() == 22
 
 
+def test_legend_labelcolor_single():
+    # test labelcolor for a single color
+    fig, ax = plt.subplots()
+    ax.plot(np.arange(10), np.arange(10)*1, label='#1')
+    ax.plot(np.arange(10), np.arange(10)*2, label='#2')
+    ax.plot(np.arange(10), np.arange(10)*3, label='#3')
+
+    leg = ax.legend(labelcolor='red')
+    for text in leg.get_texts():
+        assert mpl.colors.same_color(text.get_color(), 'red')
+
+
+def test_legend_labelcolor_list():
+    # test labelcolor for a list of colors
+    fig, ax = plt.subplots()
+    ax.plot(np.arange(10), np.arange(10)*1, label='#1')
+    ax.plot(np.arange(10), np.arange(10)*2, label='#2')
+    ax.plot(np.arange(10), np.arange(10)*3, label='#3')
+
+    leg = ax.legend(labelcolor=['r', 'g', 'b'])
+    for text, color in zip(leg.get_texts(), ['r', 'g', 'b']):
+        assert mpl.colors.same_color(text.get_color(), color)
+
+
+def test_legend_labelcolor_linecolor():
+    # test the labelcolor for labelcolor='linecolor'
+    fig, ax = plt.subplots()
+    ax.plot(np.arange(10), np.arange(10)*1, label='#1', color='r')
+    ax.plot(np.arange(10), np.arange(10)*2, label='#2', color='g')
+    ax.plot(np.arange(10), np.arange(10)*3, label='#3', color='b')
+
+    leg = ax.legend(labelcolor='linecolor')
+    for text, color in zip(leg.get_texts(), ['r', 'g', 'b']):
+        assert mpl.colors.same_color(text.get_color(), color)
+
+
+def test_legend_labelcolor_markeredgecolor():
+    # test the labelcolor for labelcolor='markeredgecolor'
+    fig, ax = plt.subplots()
+    ax.plot(np.arange(10), np.arange(10)*1, label='#1', markeredgecolor='r')
+    ax.plot(np.arange(10), np.arange(10)*2, label='#2', markeredgecolor='g')
+    ax.plot(np.arange(10), np.arange(10)*3, label='#3', markeredgecolor='b')
+
+    leg = ax.legend(labelcolor='markeredgecolor')
+    for text, color in zip(leg.get_texts(), ['r', 'g', 'b']):
+        assert mpl.colors.same_color(text.get_color(), color)
+
+
+def test_legend_labelcolor_markerfacecolor():
+    # test the labelcolor for labelcolor='markerfacecolor'
+    fig, ax = plt.subplots()
+    ax.plot(np.arange(10), np.arange(10)*1, label='#1', markerfacecolor='r')
+    ax.plot(np.arange(10), np.arange(10)*2, label='#2', markerfacecolor='g')
+    ax.plot(np.arange(10), np.arange(10)*3, label='#3', markerfacecolor='b')
+
+    leg = ax.legend(labelcolor='markerfacecolor')
+    for text, color in zip(leg.get_texts(), ['r', 'g', 'b']):
+        assert mpl.colors.same_color(text.get_color(), color)
+
+
 def test_get_set_draggable():
     legend = plt.legend()
     assert not legend.get_draggable()
