Broadcast 'orientations' arg to Sankey.add and fix its doc.

anntzer · anntzer · commit d13dfa18aec1 · 2019-01-17T23:25:51.000+01:00
diff --git a/lib/matplotlib/sankey.py b/lib/matplotlib/sankey.py
@@ -370,10 +370,9 @@ def add(self, patchlabel='', flows=None, orientations=None, labels='',
                             By convention, inputs are positive and outputs are
                             negative.
           *orientations*    list of orientations of the paths
-                            Valid values are 1 (from/to the top), 0 (from/to
-                            the left or right), or -1 (from/to the bottom).  If
-                            *orientations* == 0, inputs will break in from the
-                            left and outputs will break away to the right.
+                            Valid values are 1 (from/to the top), 0 (inputs
+                            from the left, outputs to the right), or -1
+                            (from/to the bottom).
           *labels*          list of specifications of the labels for the flows
                             Each value may be *None* (no labels), '' (just
                             label the quantities), or a labeling string.  If a
@@ -443,12 +442,14 @@ def add(self, patchlabel='', flows=None, orientations=None, labels='',
             # In the code below, angles are expressed in deg/90.
             rotation /= 90.0
         if orientations is None:
-            orientations = [0, 0]
-        if len(orientations) != n:
+            orientations = 0
+        try:
+            orientations = np.broadcast_to(orientations, n)
+        except ValueError:
             raise ValueError(
-            "orientations and flows must have the same length.\n"
-            "orientations has length %d, but flows has length %d."
-            % (len(orientations), n))
+                f"The shapes of 'flows' ({np.shape(flows)}) and "
+                f"'orientations' ({np.shape(orientations)}) are incompatible"
+            ) from None
         if not cbook.is_scalar_or_string(labels) and len(labels) != n:
             raise ValueError(
                 "If labels is a list, then labels and flows must have the "
