Add multiple labels support to Axes.plot()

yozhikoff · yozhikoff · commit 18a3bee73430 · 2020-12-23T22:56:26.000+03:00
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -1,4 +1,5 @@
 from collections import OrderedDict
+from collections.abc import Iterable
 from contextlib import ExitStack
 import functools
 import inspect
@@ -446,8 +447,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")
+
+        if ('label' in kwargs 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)}")
+
+            result = (func(x[:, j % ncx], y[:, j % ncy], kw,
+                {**kwargs, 'label':kwargs['label'][j]})
+                     for j in range(max(ncx, ncy)))
+
         result = (func(x[:, j % ncx], y[:, j % ncy], kw, kwargs)
-                  for j in range(max(ncx, ncy)))
+                for j in range(max(ncx, ncy)))
+
         if return_kwargs:
             return list(result)
         else:
