add legend to boxplot

saranti · saranti · commit 95890c2f82db · 2024-03-16T15:34:34.000+11:00
diff --git a/doc/users/next_whats_new/boxplot_legend_support.rst b/doc/users/next_whats_new/boxplot_legend_support.rst
@@ -0,0 +1,60 @@
+Legend support for Boxplot
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+Boxplots now support a *label* parameter to create legend entries.
+
+Legend labels can be passed as a list of strings to label multiple boxes in a single
+boxplot call:
+
+
+.. plot::
+    :include-source: true
+    :alt: Example of creating 3 boxplots and assigning legend labels as a sequence.
+
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    np.random.seed(19680801)
+    fruit_weights = [
+        np.random.normal(130, 10, size=100),
+        np.random.normal(125, 20, size=100),
+        np.random.normal(120, 30, size=100),
+    ]
+    labels = ['peaches', 'oranges', 'tomatoes']
+    colors = ['peachpuff', 'orange', 'tomato']
+
+    fig, ax = plt.subplots()
+    ax.set_ylabel('fruit weight (g)')
+
+    bplot = ax.boxplot(fruit_weights,
+                       patch_artist=True,  # fill with color
+                       label=labels)
+
+    # fill with colors
+    for patch, color in zip(bplot['boxes'], colors):
+        patch.set_facecolor(color)
+
+    ax.set_xticks([])
+    ax.legend()
+
+
+Or as a single string to each individual boxplot:
+
+.. plot::
+    :include-source: true
+    :alt: Example of creating 2 boxplots and assigning each legend label as a string.
+
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    fig, ax = plt.subplots()
+
+    data_A = np.random.random((100, 3))
+    data_B = np.random.random((100, 3)) + 0.2
+    pos = np.arange(3)
+
+    ax.boxplot(data_A, positions=pos - 0.2, patch_artist=True, label='Box A',
+               boxprops={'facecolor': 'steelblue'})
+    ax.boxplot(data_B, positions=pos + 0.2, patch_artist=True, label='Box B',
+               boxprops={'facecolor': 'lightblue'})
+
+    ax.legend()
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -3798,7 +3798,7 @@ def boxplot(self, x, notch=None, sym=None, vert=None, whis=None,
                 tick_labels=None, flierprops=None, medianprops=None,
                 meanprops=None, capprops=None, whiskerprops=None,
                 manage_ticks=True, autorange=False, zorder=None,
-                capwidths=None):
+                capwidths=None, label=None):
         """
         Draw a box and whisker plot.
 
@@ -3985,6 +3985,18 @@ def boxplot(self, x, notch=None, sym=None, vert=None, whis=None,
             The style of the median.
         meanprops : dict, default: None
             The style of the mean.
+        label : str or list of str, optional
+            Legend labels. Use a single string when all boxes have the same style and
+            you only want a single legend entry for them. Use a list of strings to
+            label all boxes individually. To be distinguishable, the boxes should be
+            styled individually, which is currently only possible by modifying the
+            returned artists, see e.g. :doc:`/gallery/statistics/boxplot_demo`.
+
+            In the case of a single string, the legend entry will technically be
+            associated with the first box only. By default, the legend will show the
+            median line (``result["medians"]``); if *patch_artist* is True, the legend
+            will show the box `Patch` artists (``result["boxes"]``) instead.
+
         data : indexable object, optional
             DATA_PARAMETER_PLACEHOLDER
 
@@ -4105,7 +4117,7 @@ def boxplot(self, x, notch=None, sym=None, vert=None, whis=None,
                            meanline=meanline, showfliers=showfliers,
                            capprops=capprops, whiskerprops=whiskerprops,
                            manage_ticks=manage_ticks, zorder=zorder,
-                           capwidths=capwidths)
+                           capwidths=capwidths, label=label)
         return artists
 
     def bxp(self, bxpstats, positions=None, widths=None, vert=True,
@@ -4114,7 +4126,7 @@ def bxp(self, bxpstats, positions=None, widths=None, vert=True,
             boxprops=None, whiskerprops=None, flierprops=None,
             medianprops=None, capprops=None, meanprops=None,
             meanline=False, manage_ticks=True, zorder=None,
-            capwidths=None):
+            capwidths=None, label=None):
         """
         Draw a box and whisker plot from pre-computed statistics.
 
@@ -4197,6 +4209,18 @@ def bxp(self, bxpstats, positions=None, widths=None, vert=True,
           If True, the tick locations and labels will be adjusted to match the
           boxplot positions.
 
+        label : str or list of str, optional
+            Legend labels. Use a single string when all boxes have the same style and
+            you only want a single legend entry for them. Use a list of strings to
+            label all boxes individually. To be distinguishable, the boxes should be
+            styled individually, which is currently only possible by modifying the
+            returned artists, see e.g. :doc:`/gallery/statistics/boxplot_demo`.
+
+            In the case of a single string, the legend entry will technically be
+            associated with the first box only. By default, the legend will show the
+            median line (``result["medians"]``); if *patch_artist* is True, the legend
+            will show the box `Patch` artists (``result["boxes"]``) instead.
+
         zorder : float, default: ``Line2D.zorder = 2``
           The zorder of the resulting boxplot.
 
@@ -4361,6 +4385,7 @@ def do_patch(xs, ys, **kwargs):
             if showbox:
                 do_box = do_patch if patch_artist else do_plot
                 boxes.append(do_box(box_x, box_y, **box_kw))
+                median_kw.setdefault('label', '_nolegend_')
             # draw the whiskers
             whisker_kw.setdefault('label', '_nolegend_')
             whiskers.append(do_plot(whis_x, whislo_y, **whisker_kw))
@@ -4371,7 +4396,6 @@ def do_patch(xs, ys, **kwargs):
                 caps.append(do_plot(cap_x, cap_lo, **cap_kw))
                 caps.append(do_plot(cap_x, cap_hi, **cap_kw))
             # draw the medians
-            median_kw.setdefault('label', '_nolegend_')
             medians.append(do_plot(med_x, med_y, **median_kw))
             # maybe draw the means
             if showmeans:
@@ -4389,6 +4413,19 @@ def do_patch(xs, ys, **kwargs):
                 flier_y = stats['fliers']
                 fliers.append(do_plot(flier_x, flier_y, **flier_kw))
 
+        # Set legend labels
+        if label:
+            box_or_med = boxes if showbox and patch_artist else medians
+            if cbook.is_scalar_or_string(label):
+                # assign the label only to the first box
+                box_or_med[0].set_label(label)
+            else:  # label is a sequence
+                if len(box_or_med) != len(label):
+                    raise ValueError("There must be an equal number of legend"
+                                     " labels and boxplots.")
+                for artist, lbl in zip(box_or_med, label):
+                    artist.set_label(lbl)
+
         if manage_ticks:
             axis_name = "x" if vert else "y"
             interval = getattr(self.dataLim, f"interval{axis_name}")
diff --git a/lib/matplotlib/axes/_axes.pyi b/lib/matplotlib/axes/_axes.pyi
@@ -372,6 +372,7 @@ class Axes(_AxesBase):
         autorange: bool = ...,
         zorder: float | None = ...,
         capwidths: float | ArrayLike | None = ...,
+		label: Sequence[str] | None = ...,
         *,
         data=...,
     ) -> dict[str, Any]: ...
@@ -397,6 +398,7 @@ class Axes(_AxesBase):
         manage_ticks: bool = ...,
         zorder: float | None = ...,
         capwidths: float | ArrayLike | None = ...,
+		label: Sequence[str] | None = ...,
     ) -> dict[str, Any]: ...
     def scatter(
         self,
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -2871,6 +2871,7 @@ def boxplot(
     autorange: bool = False,
     zorder: float | None = None,
     capwidths: float | ArrayLike | None = None,
+    label: Sequence[str] | None = None,
     *,
     data=None,
 ) -> dict[str, Any]:
@@ -2902,6 +2903,7 @@ def boxplot(
         autorange=autorange,
         zorder=zorder,
         capwidths=capwidths,
+        label=label,
         **({"data": data} if data is not None else {}),
     )
 
diff --git a/lib/matplotlib/tests/test_legend.py b/lib/matplotlib/tests/test_legend.py
@@ -1435,3 +1435,50 @@ def test_legend_text():
         leg_bboxes.append(
             leg.get_window_extent().transformed(ax.transAxes.inverted()))
     assert_allclose(leg_bboxes[1].bounds, leg_bboxes[0].bounds)
+
+
+def test_boxplot_labels():
+    # Test that boxplot(..., labels=) sets the tick labels but not legend entries
+    # This is not consistent with other plot types but is the current behavior.
+    fig, ax = plt.subplots()
+    np.random.seed(19680801)
+    data = np.random.random((10, 3))
+    bp = ax.boxplot(data, labels=['A', 'B', 'C'])
+    # Check that labels set the tick labels ...
+    assert [l.get_text() for l in ax.get_xticklabels()] == ['A', 'B', 'C']
+    # ... but not legend entries
+    handles, labels = ax.get_legend_handles_labels()
+    assert len(handles) == 0
+    assert len(labels) == 0
+
+
+def test_boxplot_legend_labels():
+    # Test that legend entries are generated when passing `label`.
+    np.random.seed(19680801)
+    data = np.random.random((10, 4))
+    fig, axs = plt.subplots(nrows=1, ncols=4)
+    legend_labels = ['box A', 'box B', 'box C', 'box D']
+
+    # Testing legend labels and patch passed to legend.
+    bp1 = axs[0].boxplot(data, patch_artist=True, label=legend_labels)
+    assert [v.get_label() for v in bp1['boxes']] == legend_labels
+    handles, labels = axs[0].get_legend_handles_labels()
+    assert labels == legend_labels
+    assert all(isinstance(h, mpl.patches.PathPatch) for h in handles)
+
+    # Testing legend without `box`.
+    bp2 = axs[1].boxplot(data, label=legend_labels, showbox=False)
+    # Without a box, The legend entries should be passed from the medians.
+    assert [v.get_label() for v in bp2['medians']] == legend_labels
+    handles, labels = axs[1].get_legend_handles_labels()
+    assert labels == legend_labels
+    assert all(isinstance(h, mpl.lines.Line2D) for h in handles)
+
+    # Testing legend with number of labels different from number of boxes.
+    with pytest.raises(ValueError, match='There must be an equal number'):
+        bp3 = axs[2].boxplot(data, label=legend_labels[:-1])
+
+    # Test that for a string label, only the first box gets a label.
+    bp4 = axs[3].boxplot(data, label='box A')
+    assert bp4['medians'][0].get_label() == 'box A'
+    assert all(x.get_label().startswith("_") for x in bp4['medians'][1:])
