Merge pull request #16178 from yozhikoff/add-multiple-label-support

tacaswell · web-flow · commit 6305e8da0e57 · 2020-12-24T16:31:15.000-05:00
ENH: Add multiple label support for Axes.plot()
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
@@ -0,0 +1,19 @@
+An iterable object with labels can be passed to `.Axes.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::
+
+    import matplotlib.pyplot as plt
+    
+    x = [1, 2, 3]
+
+    y = [[1, 2],
+         [2, 5],
+         [4, 9]]
+
+    plt.plot(x, y, label=['low', 'high'])
+    plt.legend()
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -1499,6 +1499,8 @@ def plot(self, *args, scalex=True, scaley=True, data=None, **kwargs):
 
             If you make multiple lines with one plot call, the kwargs
             apply to all those lines.
+            In case if label object is iterable, each its element is
+            used as label for a separate line.
 
             Here is a list of available `.Line2D` properties:
 
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -294,6 +294,12 @@ def __call__(self, *args, data=None, **kwargs):
                     replaced[label_namer_idx], args[label_namer_idx])
             args = replaced
 
+        if len(args) >= 4 and not cbook.is_scalar_or_string(
+                kwargs.get("label")):
+            raise ValueError("plot() with multiple groups of data (i.e., "
+                             "pairs of x and y) does not support multiple "
+                             "labels")
+
         # Repeatedly grab (x, y) or (x, y, format) from the front of args and
         # massage them into arguments to plot() or fill().
 
@@ -447,8 +453,22 @@ def _plot_args(self, tup, kwargs, return_kwargs=False):
         ncx, ncy = x.shape[1], y.shape[1]
         if ncx > 1 and ncy > 1 and ncx != ncy:
             raise ValueError(f"x has {ncx} columns but y has {ncy} columns")
-        result = (func(x[:, j % ncx], y[:, j % ncy], kw, kwargs)
-                  for j in range(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:
+            labels = [label] * n_datasets
+
+        result = (func(x[:, j % ncx], y[:, j % ncy], kw,
+                       {**kwargs, 'label': label})
+                  for j, label in enumerate(labels))
+
         if return_kwargs:
             return list(result)
         else:
diff --git a/lib/matplotlib/tests/test_legend.py b/lib/matplotlib/tests/test_legend.py
@@ -671,3 +671,65 @@ def test_no_warn_big_data_when_loc_specified():
         ax.plot(np.arange(5000), label=idx)
     legend = ax.legend('best')
     fig.draw_artist(legend)  # Check that no warning is emitted.
+
+
+@pytest.mark.parametrize('label_array', [['low', 'high'],
+                                         ('low', 'high'),
+                                         np.array(['low', 'high'])])
+def test_plot_multiple_input_multiple_label(label_array):
+    # test ax.plot() with multidimensional input
+    # and multiple labels
+    x = [1, 2, 3]
+    y = [[1, 2],
+         [2, 5],
+         [4, 9]]
+
+    fig, ax = plt.subplots()
+    ax.plot(x, y, label=label_array)
+    leg = ax.legend()
+    legend_texts = [entry.get_text() for entry in leg.get_texts()]
+    assert legend_texts == ['low', 'high']
+
+
+@pytest.mark.parametrize('label', ['one', 1, int])
+def test_plot_multiple_input_single_label(label):
+    # test ax.plot() with multidimensional input
+    # and single label
+    x = [1, 2, 3]
+    y = [[1, 2],
+         [2, 5],
+         [4, 9]]
+
+    fig, ax = plt.subplots()
+    ax.plot(x, y, label=label)
+    leg = ax.legend()
+    legend_texts = [entry.get_text() for entry in leg.get_texts()]
+    assert legend_texts == [str(label)] * 2
+
+
+@pytest.mark.parametrize('label_array', [['low', 'high'],
+                                         ('low', 'high'),
+                                         np.array(['low', 'high'])])
+def test_plot_single_input_multiple_label(label_array):
+    # test ax.plot() with 1D array like input
+    # and iterable label
+    x = [1, 2, 3]
+    y = [2, 5, 6]
+    fig, ax = plt.subplots()
+    ax.plot(x, y, label=label_array)
+    leg = ax.legend()
+    assert len(leg.get_texts()) == 1
+    assert leg.get_texts()[0].get_text() == str(label_array)
+
+
+def test_plot_multiple_label_incorrect_length_exception():
+    # check that excepton is raised if multiple labels
+    # are given, but number of on labels != number of lines
+    with pytest.raises(ValueError):
+        x = [1, 2, 3]
+        y = [[1, 2],
+             [2, 5],
+             [4, 9]]
+        label = ['high', 'low', 'medium']
+        fig, ax = plt.subplots()
+        ax.plot(x, y, label=label)
