Merge pull request #19187 from CharlesHe16/master

tacaswell · web-flow · commit fce534c92346 · 2021-01-04T13:43:36.000-05:00
ENH: Support callable for formatting of Sankey labels
diff --git a/doc/users/next_whats_new/callables_for_formatting_sankey_labels.rst b/doc/users/next_whats_new/callables_for_formatting_sankey_labels.rst
@@ -0,0 +1,46 @@
+Support callable for formatting of Sankey labels
+------------------------------------------------
+
+The `format` parameter of `matplotlib.sankey.Sankey` can now accept callables.
+
+This allows the use of an arbitrary function to label flows, for example allowing
+the mapping of numbers to emoji.
+
+.. plot::
+
+    import matplotlib.pyplot as plt
+    from matplotlib.sankey import Sankey
+    import math
+
+
+    def display_in_cats(values, min_cats, max_cats):
+        def display_in_cat_scale(value):
+            max_value = max(values, key=abs)
+            number_cats_to_show = \
+                max(min_cats, math.floor(abs(value) / max_value * max_cats))
+            return str(number_cats_to_show * '🐱')
+
+        return display_in_cat_scale
+
+
+    flows = [35, 15, 40, -20, -15, -5, -40, -10]
+    orientations = [-1, 1, 0, 1, 1, 1, -1, -1]
+
+    # Cats are good, we want a strictly positive number of them
+    min_cats = 1
+    # More than four cats might be too much for some people
+    max_cats = 4
+
+    cats_format = display_in_cats(flows, min_cats, max_cats)
+
+    sankey = Sankey(flows=flows, orientations=orientations, format=cats_format,
+                    offset=.1, head_angle=180, shoulder=0, scale=.010)
+
+    diagrams = sankey.finish()
+
+    diagrams[0].texts[2].set_text('')
+
+    plt.title(f'Sankey flows measured in cats \n'
+              f'🐱 = {max(flows, key=abs) / max_cats}')
+
+    plt.show()
diff --git a/lib/matplotlib/sankey.py b/lib/matplotlib/sankey.py
@@ -83,10 +83,12 @@ def __init__(self, ax=None, scale=1.0, unit='', format='%G', gap=0.25,
         unit : str
             The physical unit associated with the flow quantities.  If *unit*
             is None, then none of the quantities are labeled.
-        format : str
-            A Python number formatting string to be used in labeling the flow
-            as a quantity (i.e., a number times a unit, where the unit is
-            given).
+        format : str or callable
+            A Python number formatting string or callable used to label the
+            flows with their quantities (i.e., a number times a unit, where the
+            unit is given). If a format string is given, the label will be
+            ``format % quantity``. If a callable is given, it will be called
+            with ``quantity`` as an argument.
         gap : float
             Space between paths that break in/break away to/from the top or
             bottom.
@@ -739,7 +741,13 @@ def _get_angle(a, r):
             if label is None or angle is None:
                 label = ''
             elif self.unit is not None:
-                quantity = self.format % abs(number) + self.unit
+                if isinstance(self.format, str):
+                    quantity = self.format % abs(number) + self.unit
+                elif callable(self.format):
+                    quantity = self.format(number)
+                else:
+                    raise TypeError(
+                        'format must be callable or a format string')
                 if label != '':
                     label += "\n"
                 label += quantity
diff --git a/lib/matplotlib/tests/test_sankey.py b/lib/matplotlib/tests/test_sankey.py
@@ -10,3 +10,15 @@ def test_sankey():
 def test_label():
     s = Sankey(flows=[0.25], labels=['First'], orientations=[-1])
     assert s.diagrams[0].texts[0].get_text() == 'First\n0.25'
+
+
+def test_format_using_callable():
+    # test using callable by slightly incrementing above label example
+
+    def show_three_decimal_places(value):
+        return f'{value:.3f}'
+
+    s = Sankey(flows=[0.25], labels=['First'], orientations=[-1],
+               format=show_three_decimal_places)
+
+    assert s.diagrams[0].texts[0].get_text() == 'First\n0.250'
