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 new file mode 100644 index 000000000000..38aa766a38e4 --- /dev/null +++ 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 index 58eb69a053a9..8032001561bb 100644 --- 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 index 62c4fce662cb..1851525bd4e2 100644 --- 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'