Merge pull request #21289 from anntzer/ilb

QuLogic · web-flow · commit c1753482d52d · 2021-10-15T03:10:02.000-04:00
Simplify _init_legend_box.
diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py
@@ -712,27 +712,19 @@ def _init_legend_box(self, handles, labels, markerfirst=True):
 
         fontsize = self._fontsize
 
-        # legend_box is a HPacker, horizontally packed with
-        # columns. Each column is a VPacker, vertically packed with
-        # legend items. Each legend item is HPacker packed with
-        # legend handleBox and labelBox. handleBox is an instance of
-        # offsetbox.DrawingArea which contains legend handle. labelBox
-        # is an instance of offsetbox.TextArea which contains legend
-        # text.
+        # legend_box is a HPacker, horizontally packed with columns.
+        # Each column is a VPacker, vertically packed with legend items.
+        # Each legend item is a HPacker packed with:
+        # - handlebox: a DrawingArea which contains the legend handle.
+        # - labelbox: a TextArea which contains the legend text.
 
         text_list = []  # the list of text instances
         handle_list = []  # the list of handle instances
         handles_and_labels = []
 
-        label_prop = dict(verticalalignment='baseline',
-                          horizontalalignment='left',
-                          fontproperties=self.prop,
-                          )
-
         # The approximate height and descent of text. These values are
         # only used for plotting the legend handle.
-        descent = 0.35 * fontsize * (self.handleheight - 0.7)
-        # 0.35 and 0.7 are just heuristic numbers and may need to be improved.
+        descent = 0.35 * fontsize * (self.handleheight - 0.7)  # heuristic.
         height = fontsize * self.handleheight - descent
         # each handle needs to be drawn inside a box of (x, y, w, h) =
         # (0, -descent, width, height).  And their coordinates should
@@ -744,7 +736,7 @@ def _init_legend_box(self, handles, labels, markerfirst=True):
         # manually set their transform to the self.get_transform().
         legend_handler_map = self.get_legend_handler_map()
 
-        for orig_handle, lab in zip(handles, labels):
+        for orig_handle, label in zip(handles, labels):
             handler = self.get_legend_handler(legend_handler_map, orig_handle)
             if handler is None:
                 _api.warn_external(
@@ -753,12 +745,14 @@ def _init_legend_box(self, handles, labels, markerfirst=True):
                     "https://matplotlib.org/users/legend_guide.html"
                     "#creating-artists-specifically-for-adding-to-the-legend-"
                     "aka-proxy-artists".format(orig_handle))
-                # We don't have a handle for this artist, so we just defer
-                # to None.
+                # No handle for this artist, so we just defer to None.
                 handle_list.append(None)
             else:
-                textbox = TextArea(lab, textprops=label_prop,
-                                   multilinebaseline=True)
+                textbox = TextArea(label, multilinebaseline=True,
+                                   textprops=dict(
+                                       verticalalignment='baseline',
+                                       horizontalalignment='left',
+                                       fontproperties=self.prop))
                 handlebox = DrawingArea(width=self.handlelength * fontsize,
                                         height=height,
                                         xdescent=0., ydescent=descent)
@@ -770,34 +764,25 @@ def _init_legend_box(self, handles, labels, markerfirst=True):
                                                          fontsize, handlebox))
                 handles_and_labels.append((handlebox, textbox))
 
-        if handles_and_labels:
-            # We calculate number of rows in each column. The first
-            # (num_largecol) columns will have (nrows+1) rows, and remaining
-            # (num_smallcol) columns will have (nrows) rows.
-            ncol = min(self._ncol, len(handles_and_labels))
-            nrows, num_largecol = divmod(len(handles_and_labels), ncol)
-            num_smallcol = ncol - num_largecol
-            # starting index of each column and number of rows in it.
-            rows_per_col = [nrows + 1] * num_largecol + [nrows] * num_smallcol
-            start_idxs = np.concatenate([[0], np.cumsum(rows_per_col)[:-1]])
-            cols = zip(start_idxs, rows_per_col)
-        else:
-            cols = []
-
         columnbox = []
-        for i0, di in cols:
-            # pack handleBox and labelBox into itemBox
-            itemBoxes = [HPacker(pad=0,
+        # array_split splits n handles_and_labels into ncol columns, with the
+        # first n%ncol columns having an extra entry.  filter(len, ...) handles
+        # the case where n < ncol: the last ncol-n columns are empty and get
+        # filtered out.
+        for handles_and_labels_column \
+                in filter(len, np.array_split(handles_and_labels, self._ncol)):
+            # pack handlebox and labelbox into itembox
+            itemboxes = [HPacker(pad=0,
                                  sep=self.handletextpad * fontsize,
                                  children=[h, t] if markerfirst else [t, h],
                                  align="baseline")
-                         for h, t in handles_and_labels[i0:i0 + di]]
-            # pack columnBox
+                         for h, t in handles_and_labels_column]
+            # pack columnbox
             alignment = "baseline" if markerfirst else "right"
             columnbox.append(VPacker(pad=0,
                                      sep=self.labelspacing * fontsize,
                                      align=alignment,
-                                     children=itemBoxes))
+                                     children=itemboxes))
 
         mode = "expand" if self._mode == "expand" else "fixed"
         sep = self.columnspacing * fontsize
