Merge pull request #13218 from anntzer/sankey-labels

QuLogic · web-flow · commit 4e967f3aa700 · 2019-01-21T18:47:50.000-05:00
Fix checking of 'labels' argument to Sankey.add.
diff --git a/lib/matplotlib/sankey.py b/lib/matplotlib/sankey.py
@@ -10,7 +10,7 @@
 from matplotlib.path import Path
 from matplotlib.patches import PathPatch
 from matplotlib.transforms import Affine2D
-from matplotlib import cbook, docstring
+from matplotlib import docstring
 from matplotlib import rcParams
 
 _log = logging.getLogger(__name__)
@@ -453,12 +453,13 @@ def add(self, patchlabel='', flows=None, orientations=None, labels='',
                 f"The shapes of 'flows' {np.shape(flows)} and 'orientations' "
                 f"{np.shape(orientations)} are incompatible"
             ) from None
-        if not cbook.is_scalar_or_string(labels) and len(labels) != n:
+        try:
+            labels = np.broadcast_to(labels, n)
+        except ValueError:
             raise ValueError(
-                f"The lengths of 'flows' ({n}) and 'labels' ({len(labels)}) "
-                f"are incompatible")
-        else:
-            labels = [labels] * n
+                f"The shapes of 'flows' {np.shape(flows)} and 'labels' "
+                f"{np.shape(labels)} are incompatible"
+            ) from None
         if trunklength < 0:
             raise ValueError(
                 "'trunklength' is negative, which is not allowed because it "
diff --git a/lib/matplotlib/tests/test_sankey.py b/lib/matplotlib/tests/test_sankey.py
@@ -5,3 +5,8 @@ def test_sankey():
     # lets just create a sankey instance and check the code runs
     sankey = Sankey()
     sankey.add()
+
+
+def test_label():
+    s = Sankey(flows=[0.25], labels=['First'], orientations=[-1])
+    assert s.diagrams[0].texts[0].get_text() == 'First\n0.25'
