Apply suggestions from code review

yozhikoff · timhoffm · yozhikoff · commit 7fac6d3c304f · 2020-12-23T22:56:27.000+03:00
Co-authored-by: Tim Hoffmann &lt;2836374+timhoffm@users.noreply.github.com&gt;
diff --git a/doc/users/next_whats_new/multiple_labels_for_Axes_plot.rst b/doc/users/next_whats_new/multiple_labels_for_Axes_plot.rst
@@ -1,20 +1,17 @@
 An iterable object with labels can be passed to `.Axes.plot`
 -------------------------------------------------------------
 
-If multidimensional data is used for plotting, labels can be specified in
-a vectorized way with an iterable object of size corresponding to the
-data array shape (exactly 5 labels are expected when plotting 5 lines).
-It works with `.Axes.plot` as well as with it's wrapper `.pyplot.plot`.
+When plotting multiple datasets by passing 2D data as *y* value to `~.Axes.plot`, labels for the datasets can be passed as a list, the length matching the number of columns in *y*.
 
 .. plot::
 
-    from matplotlib import pyplot as plt
+    import matplotlib.pyplot as plt
+    
+    x = [1, 2, 3]
 
-    x = [1, 2, 5]
+    y = [[1, 9],
+         [2, 8],
+         [4, 6]]
 
-    y = [[2, 4, 3],
-        [4, 7, 1],
-        [3, 9, 2]]
-
-    plt.plot(x, y, label=['one', 'two', 'three'])
-    plt.legend()
+    plt.plot(x, y, label=['low', 'high'])
+    plt.legend()
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -448,22 +448,20 @@ def _plot_args(self, tup, kwargs, return_kwargs=False):
         if ncx > 1 and ncy > 1 and ncx != ncy:
             raise ValueError(f"x has {ncx} columns but y has {ncy} columns")
 
-        if ('label' in kwargs and max(ncx, ncy) > 1
-                and isinstance(kwargs['label'], Iterable)
-                and not isinstance(kwargs['label'], str)):
-            if len(kwargs['label']) != max(ncx, ncy):
-                raise ValueError(f"if label is iterable label and input data"
-                                 f" must have same length, but have lengths "
-                                 f"{len(kwargs['label'])} and "
-                                 f"{max(ncx, ncy)}")
-        elif 'label' in kwargs:
-            kwargs['label'] = [kwargs['label']] * max(ncx, ncy)
+        label = kwargs.get('label')
+        n_datasets = max(ncx, ncy)
+        if n_datasets > 1 and not cbook.is_scalar_or_string(label):
+            if len(label) != n_datasets:
+                raise ValueError(f"label must be scalar or have the same "
+                                 f"length as the input data, but found "
+                                 f"{len(label)} for {n_datasets} datasets.")
+            labels = label
         else:
-            kwargs['label'] = [None] * max(ncx, ncy)
+            labels = [label] * n_datasets
 
         result = (func(x[:, j % ncx], y[:, j % ncy], kw,
-                     {**kwargs, 'label': kwargs['label'][j]})
-                for j in range(max(ncx, ncy)))
+                     {**kwargs, 'label': label})
+                for j, label in enumerate(labels))
 
         if return_kwargs:
             return list(result)
diff --git a/lib/matplotlib/tests/test_legend.py b/lib/matplotlib/tests/test_legend.py
@@ -688,10 +688,8 @@ def test_plot_multiple_input_multiple_label():
         ax.plot(x, y, label=label)
         leg = ax.legend()
 
-        assert len(leg.get_texts()) == 3
-        assert leg.get_texts()[0].get_text() == 'one'
-        assert leg.get_texts()[1].get_text() == 'two'
-        assert leg.get_texts()[2].get_text() == 'three'
+        legend_texts = [entry.get_text() for entry in leg.get_texts()]
+        assert legend_texts == ['one', 'two', 'three']
 
 
 def test_plot_multiple_input_single_label():
@@ -706,10 +704,8 @@ def test_plot_multiple_input_single_label():
         ax.plot(x, y, label=label)
         leg = ax.legend()
 
-        assert len(leg.get_texts()) == 3
-        assert leg.get_texts()[0].get_text() == str(label)
-        assert leg.get_texts()[1].get_text() == str(label)
-        assert leg.get_texts()[2].get_text() == str(label)
+        legend_texts = [entry.get_text() for entry in leg.get_texts()]
+        assert legend_texts == [str(label)] * 3
 
 
 def test_plot_single_input_multiple_label():
