Add an auto-labeling helper function for bar charts

immaxchen · immaxchen · commit d82409ad0b67 · 2020-10-04T22:56:08.000+08:00
For GH #12386
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,136 @@
+"""
+==============
+Bar Label Demo
+==============
+
+This example shows how to use `blabel` helper function
+to create bar chart labels.
+
+See also the :doc:`grouped bar chart
+</gallery/lines_bars_and_markers/barchart>`,
+:doc:`stacked bar graph
+</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
+
+###############################################################################
+# Grouped bar chart
+
+fig1, ax1 = plt.subplots()
+
+rects1 = ax1.bar(ind - width/2, menMeans, width, label='Men')
+rects2 = ax1.bar(ind + width/2, womenMeans, width, label='Women')
+
+ax1.set_ylabel('Scores')
+ax1.set_title('Scores by group and gender')
+ax1.set_xticks(ind)
+ax1.set_xticklabels(('G1', 'G2', 'G3', 'G4', 'G5'))
+ax1.legend()
+
+# Basic labels
+ax1.blabel(rects1)
+ax1.blabel(rects2)
+
+plt.show()
+
+###############################################################################
+# Stacked bar plot with error bars
+
+fig2, ax2 = plt.subplots()
+
+p1 = ax2.bar(ind, menMeans, width, yerr=menStd, label='Men')
+p2 = ax2.bar(ind, womenMeans, width,
+             bottom=menMeans, yerr=womenStd, label='Women')
+
+ax2.set_ylabel('Scores')
+ax2.set_title('Scores by group and gender')
+ax2.set_xticks(ind)
+ax2.set_xticklabels(('G1', 'G2', 'G3', 'G4', 'G5'))
+ax2.legend()
+
+# Label with 'center' mode instead of the default 'edge' mode
+ax2.blabel(p1, mode='center')
+ax2.blabel(p2, mode='center')
+ax2.blabel(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))
+
+fig3, ax3 = plt.subplots()
+
+hbars1 = ax3.barh(y_pos, performance, xerr=error, align='center')
+ax3.set_yticks(y_pos)
+ax3.set_yticklabels(people)
+ax3.invert_yaxis()  # labels read top-to-bottom
+ax3.set_xlabel('Performance')
+ax3.set_title('How fast do you want to go today?')
+
+# Label with specially formatted floats
+ax3.blabel(hbars1, fmt='%.2f')
+
+plt.show()
+
+###############################################################################
+# Some of the more advanced things that one can do with bar labels
+
+fig4, ax4 = plt.subplots()
+
+hbars2 = ax4.barh(y_pos, performance, xerr=error, align='center')
+ax4.set_yticks(y_pos)
+ax4.set_yticklabels(people)
+ax4.invert_yaxis()  # labels read top-to-bottom
+ax4.set_xlabel('Performance')
+ax4.set_title('How fast do you want to go today?')
+
+# Label with given captions, custom padding, shifting and annotate option
+arrowprops = dict(color='b', arrowstyle="-|>",
+                  connectionstyle="angle,angleA=0,angleB=90,rad=20")
+
+ax4.blabel(hbars2, captions=['±%.2f' % e for e in error],
+           padding=30, shifting=20, arrowprops=arrowprops, color='b')
+
+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.blabel
+matplotlib.pyplot.blabel
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -2381,7 +2381,8 @@ def bar(self, x, height, width=0.8, bottom=None, *, align="center",
 
         self._request_autoscale_view()
 
-        bar_container = BarContainer(patches, errorbar, label=label)
+        bar_container = BarContainer(patches, errorbar, orientation,
+                                     label=label)
         self.add_container(bar_container)
 
         if tick_labels is not None:
@@ -2497,6 +2498,136 @@ def barh(self, y, width, height=0.8, left=None, *, align="center",
                            align=align, **kwargs)
         return patches
 
+    def blabel(self, container, captions=[], fmt="%g", mode="edge",
+               padding=None, shifting=0, autoscale=True, **kwargs):
+        """
+        Label a bar plot.
+
+        Adds labels to bars in the given `BarContainer`.
+
+        Parameters
+        ----------
+        container : `.BarContainer`
+            Container with all the bars and optionally errorbars.
+
+        captions : 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, optional
+            A format string for the label. Default is '%g'
+
+        mode : {'edge', 'center'}, optional, default: 'edge'
+            Position of the label relative to the bar:
+
+            - 'edge': Placed at edge, the cumulative values will be shown.
+            - 'center': Placed at center, the individual values will be shown.
+
+        padding : float, optional
+            Space in points for label to leave the edge of the bar.
+
+        shifting : float, optional
+            Translation in points for label to leave the center of the bar.
+
+        autoscale : bool, optional
+            If ``True``, try to rescale to fit the labels. Default is ``True``.
+
+        Returns
+        -------
+        annotations
+            A list of `.Text` instances for the labels.
+        """
+        def nonefill(a, n):
+            return list(a) + [None] * (n - len(a))
+
+        def sign(x):
+            return 1 if x >= 0 else -1
+
+        if mode == "edge":
+            padding = padding or 3
+        elif mode == "center":
+            padding = padding or 0
+
+        bars = container.patches
+        errorbar = container.errorbar
+        orientation = container.orientation
+
+        N = len(bars)
+        if errorbar:
+            lines = errorbar.lines
+            barlinecols = lines[2]
+            barlinecol = barlinecols[0]
+            segments = barlinecol.get_segments()
+            errs = nonefill(segments, N)
+        else:
+            errs = nonefill([], N)
+        caps = nonefill(captions, N)
+
+        annotations = []
+        for bar, err, cap in zip(bars, errs, caps):
+
+            (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 yc >= 0 else min(y0, y1)
+                length = abs(y0 - y1)
+            elif orientation == "horizontal":
+                extrema = max(x0, x1) if xc >= 0 else min(x0, x1)
+                length = abs(x0 - x1)
+
+            if err is None:
+                endpt = extrema
+            elif orientation == "vertical":
+                endpt = err[:, 1].max() if yc >= 0 else err[:, 1].min()
+            elif orientation == "horizontal":
+                endpt = err[:, 0].max() if xc >= 0 else err[:, 0].min()
+
+            if mode == "center":
+                value = sign(extrema) * length
+            elif mode == "edge":
+                value = extrema
+
+            if mode == "center":
+                xy = xc, yc
+            elif mode == "edge" and orientation == "vertical":
+                xy = xc, endpt
+            elif mode == "edge" and orientation == "horizontal":
+                xy = endpt, yc
+
+            if orientation == "vertical":
+                xytext = shifting, sign(extrema) * padding
+            else:
+                xytext = sign(extrema) * padding, shifting
+
+            if mode == "center":
+                ha, va = "center", "center"
+            elif mode == "edge" and orientation == "vertical" and yc >= 0:
+                ha, va = "center", "bottom"
+            elif mode == "edge" and orientation == "vertical" and yc < 0:
+                ha, va = "center", "top"
+            elif mode == "edge" and orientation == "horizontal" and xc >= 0:
+                ha, va = "left", "center"
+            elif mode == "edge" and orientation == "horizontal" and xc < 0:
+                ha, va = "right", "center"
+
+            annotation = self.annotate(cap or fmt % value, xy, xytext,
+                                       textcoords="offset points",
+                                       ha=ha, va=va, **kwargs)
+            annotations.append(annotation)
+
+        if autoscale:
+            transform = self.transData.inverted()
+            renderer = self.figure.canvas.get_renderer()
+            corners = []
+            for text in annotations:
+                points = text.get_window_extent(renderer).get_points()
+                corners.append(transform.transform(points))
+            self.update_datalim(np.vstack(corners))
+            self.autoscale_view()
+
+        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
@@ -62,9 +62,10 @@ class BarContainer(Container):
 
     """
 
-    def __init__(self, patches, errorbar=None, **kwargs):
+    def __init__(self, patches, errorbar=None, orientation=None, **kwargs):
         self.patches = patches
         self.errorbar = errorbar
+        self.orientation = orientation
         super().__init__(patches, **kwargs)
 
 
diff --git a/tools/boilerplate.py b/tools/boilerplate.py
@@ -199,6 +199,7 @@ def boilerplate_gen():
         'bar',
         'barbs',
         'barh',
+        'blabel',
         'boxplot',
         'broken_barh',
         'cla',
