Merge pull request #15602 from immaxchen/bar-chart-auto-label-gh12386

QuLogic · web-flow · commit 6a322de789b5 · 2020-12-22T04:38:25.000-05:00
Add an auto-labeling helper function for bar charts
diff --git a/doc/api/axes_api.rst b/doc/api/axes_api.rst
@@ -66,6 +66,7 @@ Basic
 
    Axes.bar
    Axes.barh
+   Axes.bar_label
 
    Axes.stem
    Axes.eventplot
diff --git a/doc/users/next_whats_new/2019-11-06_auto-labeling-for-bar-charts.rst b/doc/users/next_whats_new/2019-11-06_auto-labeling-for-bar-charts.rst
@@ -0,0 +1,7 @@
+:orphan:
+
+Bar charts auto-labeling
+------------------------
+A new `.Axes.bar_label` method has been added for auto-labeling bar charts.
+See :doc:`/gallery/lines_bars_and_markers/bar_label_demo` for examples.
+
diff --git a/doc/users/next_whats_new/2019-11-12_bar_container_accepts_orientation.rst b/doc/users/next_whats_new/2019-11-12_bar_container_accepts_orientation.rst
@@ -0,0 +1,7 @@
+:orphan:
+
+Setting BarContainer orientation
+--------------------------------
+`.BarContainer` now accepts a new string argument ``orientation``.
+It can be either ``vertical`` or ``horizontal``, default is ``None``.
+
diff --git a/examples/lines_bars_and_markers/bar_label_demo.py b/examples/lines_bars_and_markers/bar_label_demo.py
@@ -0,0 +1,116 @@
+"""
+==============
+Bar Label Demo
+==============
+
+This example shows how to use the `~.Axes.bar_label` helper function
+to create bar chart labels.
+
+See also the :doc:`grouped bar
+</gallery/lines_bars_and_markers/barchart>`,
+:doc:`stacked bar
+</gallery/lines_bars_and_markers/bar_stacked>` and
+:doc:`horizontal bar chart
+</gallery/lines_bars_and_markers/barh>` examples.
+"""
+
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+
+###############################################################################
+# Define the data
+
+N = 5
+menMeans = (20, 35, 30, 35, -27)
+womenMeans = (25, 32, 34, 20, -25)
+menStd = (2, 3, 4, 1, 2)
+womenStd = (3, 5, 2, 3, 3)
+ind = np.arange(N)    # the x locations for the groups
+width = 0.35       # the width of the bars: can also be len(x) sequence
+
+###############################################################################
+# Stacked bar plot with error bars
+
+fig, ax = plt.subplots()
+
+p1 = ax.bar(ind, menMeans, width, yerr=menStd, label='Men')
+p2 = ax.bar(ind, womenMeans, width,
+            bottom=menMeans, yerr=womenStd, label='Women')
+
+ax.axhline(0, color='grey', linewidth=0.8)
+ax.set_ylabel('Scores')
+ax.set_title('Scores by group and gender')
+ax.set_xticks(ind)
+ax.set_xticklabels(('G1', 'G2', 'G3', 'G4', 'G5'))
+ax.legend()
+
+# Label with label_type 'center' instead of the default 'edge'
+ax.bar_label(p1, label_type='center')
+ax.bar_label(p2, label_type='center')
+ax.bar_label(p2)
+
+plt.show()
+
+###############################################################################
+# Horizontal bar chart
+
+# Fixing random state for reproducibility
+np.random.seed(19680801)
+
+# Example data
+people = ('Tom', 'Dick', 'Harry', 'Slim', 'Jim')
+y_pos = np.arange(len(people))
+performance = 3 + 10 * np.random.rand(len(people))
+error = np.random.rand(len(people))
+
+fig, ax = plt.subplots()
+
+hbars = ax.barh(y_pos, performance, xerr=error, align='center')
+ax.set_yticks(y_pos)
+ax.set_yticklabels(people)
+ax.invert_yaxis()  # labels read top-to-bottom
+ax.set_xlabel('Performance')
+ax.set_title('How fast do you want to go today?')
+
+# Label with specially formatted floats
+ax.bar_label(hbars, fmt='%.2f')
+ax.set_xlim(right=15)  # adjust xlim to fit labels
+
+plt.show()
+
+###############################################################################
+# Some of the more advanced things that one can do with bar labels
+
+fig, ax = plt.subplots()
+
+hbars = ax.barh(y_pos, performance, xerr=error, align='center')
+ax.set_yticks(y_pos)
+ax.set_yticklabels(people)
+ax.invert_yaxis()  # labels read top-to-bottom
+ax.set_xlabel('Performance')
+ax.set_title('How fast do you want to go today?')
+
+# Label with given captions, custom padding and annotate options
+ax.bar_label(hbars, labels=['±%.2f' % e for e in error],
+             padding=8, color='b', fontsize=14)
+ax.set_xlim(right=16)
+
+plt.show()
+
+#############################################################################
+#
+# ------------
+#
+# References
+# """"""""""
+#
+# The use of the following functions, methods and classes is shown
+# in this example:
+
+matplotlib.axes.Axes.bar
+matplotlib.pyplot.bar
+matplotlib.axes.Axes.barh
+matplotlib.pyplot.barh
+matplotlib.axes.Axes.bar_label
+matplotlib.pyplot.bar_label
diff --git a/examples/lines_bars_and_markers/barchart.py b/examples/lines_bars_and_markers/barchart.py
@@ -30,20 +30,8 @@
 ax.set_xticklabels(labels)
 ax.legend()
 
-
-def autolabel(rects):
-    """Attach a text label above each bar in *rects*, displaying its height."""
-    for rect in rects:
-        height = rect.get_height()
-        ax.annotate('{}'.format(height),
-                    xy=(rect.get_x() + rect.get_width() / 2, height),
-                    xytext=(0, 3),  # 3 points vertical offset
-                    textcoords="offset points",
-                    ha='center', va='bottom')
-
-
-autolabel(rects1)
-autolabel(rects2)
+ax.bar_label(rects1, padding=3)
+ax.bar_label(rects2, padding=3)
 
 fig.tight_layout()
 
@@ -61,5 +49,5 @@ def autolabel(rects):
 
 matplotlib.axes.Axes.bar
 matplotlib.pyplot.bar
-matplotlib.axes.Axes.annotate
-matplotlib.pyplot.annotate
+matplotlib.axes.Axes.bar_label
+matplotlib.pyplot.bar_label
diff --git a/examples/lines_bars_and_markers/horizontal_barchart_distribution.py b/examples/lines_bars_and_markers/horizontal_barchart_distribution.py
@@ -54,15 +54,12 @@ def survey(results, category_names):
     for i, (colname, color) in enumerate(zip(category_names, category_colors)):
         widths = data[:, i]
         starts = data_cum[:, i] - widths
-        ax.barh(labels, widths, left=starts, height=0.5,
-                label=colname, color=color)
-        xcenters = starts + widths / 2
+        rects = ax.barh(labels, widths, left=starts, height=0.5,
+                        label=colname, color=color)
 
         r, g, b, _ = color
         text_color = 'white' if r * g * b < 0.5 else 'darkgrey'
-        for y, (x, c) in enumerate(zip(xcenters, widths)):
-            ax.text(x, y, str(int(c)), ha='center', va='center',
-                    color=text_color)
+        ax.bar_label(rects, label_type='center', color=text_color)
     ax.legend(ncol=len(category_names), bbox_to_anchor=(0, 1),
               loc='lower left', fontsize='small')
 
@@ -85,7 +82,7 @@ def survey(results, category_names):
 import matplotlib
 matplotlib.axes.Axes.barh
 matplotlib.pyplot.barh
-matplotlib.axes.Axes.text
-matplotlib.pyplot.text
+matplotlib.axes.Axes.bar_label
+matplotlib.pyplot.bar_label
 matplotlib.axes.Axes.legend
 matplotlib.pyplot.legend
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -2385,7 +2385,13 @@ def bar(self, x, height, width=0.8, bottom=None, *, align="center",
 
         self._request_autoscale_view()
 
-        bar_container = BarContainer(patches, errorbar, label=label)
+        if orientation == 'vertical':
+            datavalues = height
+        elif orientation == 'horizontal':
+            datavalues = width
+
+        bar_container = BarContainer(patches, errorbar, datavalues=datavalues,
+                                     orientation=orientation, label=label)
         self.add_container(bar_container)
 
         if tick_labels is not None:
@@ -2501,6 +2507,132 @@ def barh(self, y, width, height=0.8, left=None, *, align="center",
                            align=align, **kwargs)
         return patches
 
+    def bar_label(self, container, labels=None, *, fmt="%g", label_type="edge",
+                  padding=0, **kwargs):
+        """
+        Label a bar plot.
+
+        Adds labels to bars in the given `.BarContainer`.
+        You may need to adjust the axis limits to fit the labels.
+
+        Parameters
+        ----------
+        container : `.BarContainer`
+            Container with all the bars and optionally errorbars, likely
+            returned from `.bar` or `.barh`.
+
+        labels : array-like, optional
+            A list of label texts, that should be displayed. If not given, the
+            label texts will be the data values formatted with *fmt*.
+
+        fmt : str, default: '%g'
+            A format string for the label.
+
+        label_type : {'edge', 'center'}, default: 'edge'
+            The label type. Possible values:
+
+            - 'edge': label placed at the end-point of the bar segment, and the
+              value displayed will be the position of that end-point.
+            - 'center': label placed in the center of the bar segment, and the
+              value displayed will be the length of that segment.
+              (useful for stacked bars, i.e.
+              :doc:`/gallery/lines_bars_and_markers/bar_label_demo`)
+
+        padding : float, default: 0
+            Distance of label from the end of the bar.
+
+        **kwargs
+            Any remaining keyword arguments are passed through to
+            `.Axes.annotate`.
+
+        Returns
+        -------
+        list of `.Text`
+            A list of `.Text` instances for the labels.
+        """
+
+        # want to know whether to put label on positive or negative direction
+        # cannot use np.sign here because it will return 0 if x == 0
+        def sign(x):
+            return 1 if x >= 0 else -1
+
+        _api.check_in_list(['edge', 'center'], label_type=label_type)
+
+        bars = container.patches
+        errorbar = container.errorbar
+        datavalues = container.datavalues
+        orientation = container.orientation
+
+        if errorbar:
+            # check "ErrorbarContainer" for the definition of these elements
+            lines = errorbar.lines  # attribute of "ErrorbarContainer" (tuple)
+            barlinecols = lines[2]  # 0: data_line, 1: caplines, 2: barlinecols
+            barlinecol = barlinecols[0]  # the "LineCollection" of error bars
+            errs = barlinecol.get_segments()
+        else:
+            errs = []
+
+        if labels is None:
+            labels = []
+
+        annotations = []
+
+        for bar, err, dat, lbl in itertools.zip_longest(
+            bars, errs, datavalues, labels
+        ):
+            (x0, y0), (x1, y1) = bar.get_bbox().get_points()
+            xc, yc = (x0 + x1) / 2, (y0 + y1) / 2
+
+            if orientation == "vertical":
+                extrema = max(y0, y1) if dat >= 0 else min(y0, y1)
+                length = abs(y0 - y1)
+            elif orientation == "horizontal":
+                extrema = max(x0, x1) if dat >= 0 else min(x0, x1)
+                length = abs(x0 - x1)
+
+            if err is None:
+                endpt = extrema
+            elif orientation == "vertical":
+                endpt = err[:, 1].max() if dat >= 0 else err[:, 1].min()
+            elif orientation == "horizontal":
+                endpt = err[:, 0].max() if dat >= 0 else err[:, 0].min()
+
+            if label_type == "center":
+                value = sign(dat) * length
+            elif label_type == "edge":
+                value = extrema
+
+            if label_type == "center":
+                xy = xc, yc
+            elif label_type == "edge" and orientation == "vertical":
+                xy = xc, endpt
+            elif label_type == "edge" and orientation == "horizontal":
+                xy = endpt, yc
+
+            if orientation == "vertical":
+                xytext = 0, sign(dat) * padding
+            else:
+                xytext = sign(dat) * padding, 0
+
+            if label_type == "center":
+                ha, va = "center", "center"
+            elif label_type == "edge":
+                if orientation == "vertical" and dat >= 0:
+                    ha, va = "center", "bottom"
+                elif orientation == "vertical" and dat < 0:
+                    ha, va = "center", "top"
+                elif orientation == "horizontal" and dat >= 0:
+                    ha, va = "left", "center"
+                elif orientation == "horizontal" and dat < 0:
+                    ha, va = "right", "center"
+
+            annotation = self.annotate(fmt % value if lbl is None else lbl,
+                                       xy, xytext, textcoords="offset points",
+                                       ha=ha, va=va, **kwargs)
+            annotations.append(annotation)
+
+        return annotations
+
     @_preprocess_data()
     @docstring.dedent_interpd
     def broken_barh(self, xranges, yrange, **kwargs):
diff --git a/lib/matplotlib/container.py b/lib/matplotlib/container.py
@@ -57,11 +57,21 @@ class BarContainer(Container):
         A container for the error bar artists if error bars are present.
         *None* otherwise.
 
+    datavalues : None or array-like
+        The underlying data values corresponding to the bars.
+
+    orientation : {'vertical', 'horizontal'}, default: None
+        If 'vertical', the bars are assumed to be vertical.
+        If 'horizontal', the bars are assumed to be horizontal.
+
     """
 
-    def __init__(self, patches, errorbar=None, **kwargs):
+    def __init__(self, patches, errorbar=None, *, datavalues=None,
+                 orientation=None, **kwargs):
         self.patches = patches
         self.errorbar = errorbar
+        self.datavalues = datavalues
+        self.orientation = orientation
         super().__init__(patches, **kwargs)
 
 
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -2596,6 +2596,16 @@ def barh(y, width, height=0.8, left=None, *, align='center', **kwargs):
         y, width, height=height, left=left, align=align, **kwargs)
 
 
+# Autogenerated by boilerplate.py.  Do not edit as changes will be lost.
+@_copy_docstring_and_deprecators(Axes.bar_label)
+def bar_label(
+        container, labels=None, *, fmt='%g', label_type='edge',
+        padding=0, **kwargs):
+    return gca().bar_label(
+        container, labels=labels, fmt=fmt, label_type=label_type,
+        padding=padding, **kwargs)
+
+
 # Autogenerated by boilerplate.py.  Do not edit as changes will be lost.
 @_copy_docstring_and_deprecators(Axes.boxplot)
 def boxplot(
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
diff --git a/tools/boilerplate.py b/tools/boilerplate.py
