diff --git a/doc/users/next_whats_new/more_legend_handlers.rst b/doc/users/next_whats_new/more_legend_handlers.rst new file mode 100644 index 000000000000..1b6051933939 --- /dev/null +++ b/doc/users/next_whats_new/more_legend_handlers.rst @@ -0,0 +1,5 @@ +Add legend handlers for FancyArrowPatch, Text, and Annotation +------------------------------------------------------------- + +By setting the label on FancyArrowPatch, Text, and Annotation objects +they will automatically be included in legends. diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py index a95080ddea1a..e32a34337248 100644 --- a/lib/matplotlib/legend.py +++ b/lib/matplotlib/legend.py @@ -39,7 +39,8 @@ import matplotlib.colors as colors from matplotlib.font_manager import FontProperties from matplotlib.lines import Line2D -from matplotlib.patches import Patch, Rectangle, Shadow, FancyBboxPatch +from matplotlib.patches import (Patch, Rectangle, Shadow, FancyBboxPatch, + FancyArrowPatch) from matplotlib.collections import (LineCollection, RegularPolyCollection, CircleCollection, PathCollection, PolyCollection) @@ -50,6 +51,7 @@ from matplotlib.offsetbox import DraggableOffsetBox from matplotlib.container import ErrorbarContainer, BarContainer, StemContainer +from matplotlib.text import Annotation, Text from . import legend_handler @@ -565,7 +567,6 @@ def _set_loc(self, loc): # value of the find_offset. self._loc_real = loc self.stale = True - self._legend_box.set_offset(self._findoffset) def _get_loc(self): return self._loc_real @@ -649,7 +650,10 @@ def _approx_text_height(self, renderer=None): update_func=legend_handler.update_from_first_child), tuple: legend_handler.HandlerTuple(), PathCollection: legend_handler.HandlerPathCollection(), - PolyCollection: legend_handler.HandlerPolyCollection() + PolyCollection: legend_handler.HandlerPolyCollection(), + FancyArrowPatch: legend_handler.HandlerFancyArrowPatch(), + Text: legend_handler.HandlerText(), + Annotation: legend_handler.HandlerAnnotation() } # (get|set|update)_default_handler_maps are public interfaces to @@ -837,6 +841,7 @@ def _init_legend_box(self, handles, labels, markerfirst=True): children=[self._legend_title_box, self._legend_handle_box]) self._legend_box.set_figure(self.figure) + self._legend_box.set_offset(self._findoffset) self.texts = text_list self.legendHandles = handle_list @@ -1142,12 +1147,13 @@ def _get_legend_handles(axs, legend_handler_map=None): handles_original = [] for ax in axs: handles_original += (ax.lines + ax.patches + - ax.collections + ax.containers) + ax.collections + ax.containers + ax.texts) # support parasite axes: if hasattr(ax, 'parasites'): for axx in ax.parasites: handles_original += (axx.lines + axx.patches + - axx.collections + axx.containers) + axx.collections + axx.containers + + axx.texts) handler_map = Legend.get_default_handler_map() diff --git a/lib/matplotlib/legend_handler.py b/lib/matplotlib/legend_handler.py index 0968a5c4b99b..a1b6a9bd7cd3 100644 --- a/lib/matplotlib/legend_handler.py +++ b/lib/matplotlib/legend_handler.py @@ -33,8 +33,9 @@ def legend_artist(self, legend, orig_handle, fontsize, handlebox): import numpy as np from matplotlib.lines import Line2D -from matplotlib.patches import Rectangle +from matplotlib.patches import Rectangle, FancyArrowPatch import matplotlib.collections as mcoll +from matplotlib.text import Text, Annotation import matplotlib.colors as mcolors @@ -58,6 +59,7 @@ def create_artists(self, legend, orig_handle, width, height) that are scaled by fontsize if necessary. """ + def __init__(self, xpad=0., ypad=0., update_func=None): self._xpad, self._ypad = xpad, ypad self._update_prop_func = update_func @@ -110,10 +112,10 @@ def legend_artist(self, legend, orig_handle, """ xdescent, ydescent, width, height = self.adjust_drawing_area( - legend, orig_handle, - handlebox.xdescent, handlebox.ydescent, - handlebox.width, handlebox.height, - fontsize) + legend, orig_handle, + handlebox.xdescent, handlebox.ydescent, + handlebox.width, handlebox.height, + fontsize) artists = self.create_artists(legend, orig_handle, xdescent, ydescent, width, height, fontsize, handlebox.get_transform()) @@ -135,6 +137,7 @@ class HandlerNpoints(HandlerBase): """ A legend handler that shows *numpoints* points in the legend entry. """ + def __init__(self, marker_pad=0.3, numpoints=None, **kw): """ Parameters @@ -180,6 +183,7 @@ class HandlerNpointsYoffsets(HandlerNpoints): A legend handler that shows *numpoints* in the legend, and allows them to be individually offest in the y-direction. """ + def __init__(self, numpoints=None, yoffsets=None, **kw): """ Parameters @@ -211,6 +215,7 @@ class HandlerLine2D(HandlerNpoints): """ Handler for `.Line2D` instances. """ + def __init__(self, marker_pad=0.3, numpoints=None, **kw): """ Parameters @@ -263,6 +268,7 @@ class HandlerPatch(HandlerBase): """ Handler for `.Patch` instances. """ + def __init__(self, patch_func=None, **kw): """ Parameters @@ -309,6 +315,7 @@ class HandlerLineCollection(HandlerLine2D): """ Handler for `.LineCollection` instances. """ + def get_numpoints(self, legend): if self._numpoints is None: return legend.scatterpoints @@ -341,6 +348,7 @@ class HandlerRegularPolyCollection(HandlerNpointsYoffsets): """ Handler for `.RegularPolyCollections`. """ + def __init__(self, yoffsets=None, sizes=None, **kw): HandlerNpointsYoffsets.__init__(self, yoffsets=yoffsets, **kw) @@ -378,7 +386,7 @@ def update_prop(self, legend_handle, orig_handle, legend): self._update_prop(legend_handle, orig_handle) legend_handle.set_figure(legend.figure) - #legend._set_artist_props(legend_handle) + # legend._set_artist_props(legend_handle) legend_handle.set_clip_box(None) legend_handle.set_clip_path(None) @@ -416,6 +424,7 @@ class HandlerPathCollection(HandlerRegularPolyCollection): """ Handler for `.PathCollections`, which are used by `~.Axes.scatter`. """ + def create_collection(self, orig_handle, sizes, offsets, transOffset): p = type(orig_handle)([orig_handle.get_paths()[0]], sizes=sizes, @@ -429,6 +438,7 @@ class HandlerCircleCollection(HandlerRegularPolyCollection): """ Handler for `.CircleCollections`. """ + def create_collection(self, orig_handle, sizes, offsets, transOffset): p = type(orig_handle)(sizes, offsets=offsets, @@ -441,6 +451,7 @@ class HandlerErrorbar(HandlerLine2D): """ Handler for Errorbars. """ + def __init__(self, xerr_size=0.5, yerr_size=None, marker_pad=0.3, numpoints=None, **kw): @@ -503,8 +514,8 @@ def create_artists(self, legend, orig_handle, handle_caplines = [] if orig_handle.has_xerr: - verts = [ ((x - xerr_size, y), (x + xerr_size, y)) - for x, y in zip(xdata_marker, ydata_marker)] + verts = [((x - xerr_size, y), (x + xerr_size, y)) + for x, y in zip(xdata_marker, ydata_marker)] coll = mcoll.LineCollection(verts) self.update_prop(coll, barlinecols[0], legend) handle_barlinecols.append(coll) @@ -521,8 +532,8 @@ def create_artists(self, legend, orig_handle, handle_caplines.append(capline_right) if orig_handle.has_yerr: - verts = [ ((x, y - yerr_size), (x, y + yerr_size)) - for x, y in zip(xdata_marker, ydata_marker)] + verts = [((x, y - yerr_size), (x, y + yerr_size)) + for x, y in zip(xdata_marker, ydata_marker)] coll = mcoll.LineCollection(verts) self.update_prop(coll, barlinecols[0], legend) handle_barlinecols.append(coll) @@ -554,6 +565,7 @@ class HandlerStem(HandlerNpointsYoffsets): """ Handler for plots produced by `~.Axes.stem`. """ + def __init__(self, marker_pad=0.3, numpoints=None, bottom=None, yoffsets=None, **kw): """ @@ -649,11 +661,37 @@ class HandlerTuple(HandlerBase): pad : float, optional If None, fall back to ``legend.borderpad`` as the default. In units of fraction of font size. Default is None. + + + width_ratios : tuple, optional + Specifies the respective widths of a text/arrow legend annotation pair. + Must be of length ndivide. + If None, all sections will have the same width. Default is None. + + + handlers : tuple, optionnal + The list of handlers to call for each text/arrow legend annotation section. + Must be of length ndivide. + If None, the default handlers will be fetched automatically. Default is None. """ - def __init__(self, ndivide=1, pad=None, **kwargs): + + def __init__( + self, + ndivide=1, + pad=None, + width_ratios=None, + handlers=None, + **kwargs): self._ndivide = ndivide self._pad = pad + self._handlers = handlers + + if (width_ratios is not None) and (len(width_ratios) == ndivide): + self._width_ratios = width_ratios + else: + self._width_ratios = None + HandlerBase.__init__(self, **kwargs) def create_artists(self, legend, orig_handle, @@ -672,17 +710,32 @@ def create_artists(self, legend, orig_handle, else: pad = self._pad * fontsize - if ndivide > 1: - width = (width - pad * (ndivide - 1)) / ndivide + if self._width_ratios is not None: + sumratios = sum(self._width_ratios) + widths = [(width - pad * (ndivide - 1)) * ratio / sumratios + for ratio in self._width_ratios] + else: + widths = [(width - pad * (ndivide - 1)) / ndivide + for _ in range(ndivide)] + widths_cycle = cycle(widths) - xds_cycle = cycle(xdescent - (width + pad) * np.arange(ndivide)) + xds = [xdescent - (widths[-i - 1] + pad) * i for i in range(ndivide)] + xds_cycle = cycle(xds) a_list = [] - for handle1 in orig_handle: - handler = legend.get_legend_handler(handler_map, handle1) - _a_list = handler.create_artists( - legend, handle1, - next(xds_cycle), ydescent, width, height, fontsize, trans) + for i, handle1 in enumerate(orig_handle): + if self._handlers is not None: + handler = self._handlers[i] + else: + handler = legend.get_legend_handler(handler_map, handle1) + + _a_list = handler.create_artists(legend, handle1, + next(xds_cycle), + ydescent, + next(widths_cycle), + height, + fontsize, + trans) a_list.extend(_a_list) return a_list @@ -692,6 +745,7 @@ class HandlerPolyCollection(HandlerBase): """ Handler for `.PolyCollection` used in `~.Axes.fill_between` and `~.Axes.stackplot`. """ + def _update_prop(self, legend_handle, orig_handle): def first_color(colors): if colors is None: @@ -728,3 +782,164 @@ def create_artists(self, legend, orig_handle, self.update_prop(p, orig_handle, legend) p.set_transform(trans) return [p] + + +class HandlerFancyArrowPatch(HandlerPatch): + """ + Handler for FancyArrowPatch instances. + """ + + def _create_patch(self, legend, orig_handle, + xdescent, ydescent, width, height, fontsize): + arrow = FancyArrowPatch([-xdescent, + -ydescent + height / 2], + [-xdescent + width, + -ydescent + height / 2], + mutation_scale=width / 3) + arrow.set_arrowstyle(orig_handle.get_arrowstyle()) + return arrow + + +class HandlerText(HandlerBase): + """ + Handler for Text instances. + + Additional kwargs are passed through to `HandlerBase`. + + Parameters + ---------- + rep_str : string, optional + Replacement string used in the legend when the Text string is longer than rep_maxlen. + Default is 'Aa'. + + + rep_maxlen : int, optional + Maximum length of Text string to be used in the legend. Default is 2. + """ + + def __init__(self, rep_str='Aa', rep_maxlen=2, **kwargs): + + self._rep_str = rep_str + self._rep_maxlen = rep_maxlen + + HandlerBase.__init__(self, **kwargs) + + def create_artists(self, legend, orig_handle, + xdescent, ydescent, width, height, fontsize, trans): + # Use original text if it is short + text = orig_handle.get_text() + if len(text) > self._rep_maxlen: + text = self._rep_str + + # Use smaller fontsize for text repr + text_fontsize = 2 * fontsize / 3 + + t = Text(x=-xdescent + width / 2 - len(text) * text_fontsize / 4, + y=-ydescent + height / 4, + text=text) + + # Copy text attributes, except fontsize + self.update_prop(t, orig_handle, legend) + t.set_transform(trans) + t.set_fontsize(text_fontsize) + + return [t] + + +class HandlerAnnotation(HandlerText): + """ + Handler for Annotation instances. + + Defers to HandlerText to draw the annotation text (if any). + Defers to HandlerFancyArrowPatch to draw the annotation arrow (if any). + For annotations made of both text and arrow, HandlerTuple is used to draw them side by side. + Additional kwargs are passed through to `HandlerText`. + + Parameters + ---------- + + pad : float, optional + If None, fall back to `legend.borderpad` asstr the default. + In units of fraction of font size. + Default is None. + + width_ratios : tuple, optional + The relative width of the respective text/arrow legend annotation pair. + Must be of length 2. + Default is [1,4]. + """ + + def __init__( + self, + pad=None, + width_ratios=[1, 4], + **kwargs + ): + + self._pad = pad + self._width_ratios = width_ratios + + HandlerText.__init__(self, **kwargs) + + def create_artists( + self, + legend, + orig_handle, + xdescent, + ydescent, + width, + height, + fontsize, + trans, + ): + if orig_handle.arrow_patch is not None \ + and orig_handle.get_text() is not '': + + # Draw a tuple (text, arrow) + + handler = HandlerTuple( + ndivide=2, + pad=self._pad, + width_ratios=self._width_ratios, + handlers=[ + HandlerText( + rep_str=self._rep_str, + rep_maxlen=self._rep_maxlen), + HandlerFancyArrowPatch()]) + + # Create a Text instance from annotation text + + text_handle = Text(text=orig_handle.get_text()) + text_handle.update_from(orig_handle) + handle = (text_handle, orig_handle.arrow_patch) + elif orig_handle.arrow_patch is not None: + + # Arrow without text + + handler = HandlerFancyArrowPatch() + handle = orig_handle.arrow_patch + elif orig_handle.get_text() is not '': + + # Text without arrow + + handler = HandlerText(rep_str=self._rep_str, + rep_maxlen=self._rep_maxlen) + handle = orig_handle + else: + + # No text, no arrow + + handler = HandlerPatch() + handle = Rectangle(xy=[0, 0], width=0, height=0, color='w', + alpha=0.0) + + return handler.create_artists( + legend, + handle, + xdescent, + ydescent, + width, + height, + fontsize, + trans, + ) diff --git a/lib/matplotlib/tests/baseline_images/test_legend/legend_all_annotation_text.png b/lib/matplotlib/tests/baseline_images/test_legend/legend_all_annotation_text.png new file mode 100644 index 000000000000..d49cd37ea313 Binary files /dev/null and b/lib/matplotlib/tests/baseline_images/test_legend/legend_all_annotation_text.png differ diff --git a/lib/matplotlib/tests/test_legend.py b/lib/matplotlib/tests/test_legend.py index 830fe798c44c..8c3232d38b9f 100644 --- a/lib/matplotlib/tests/test_legend.py +++ b/lib/matplotlib/tests/test_legend.py @@ -10,7 +10,8 @@ import matplotlib as mpl import matplotlib.transforms as mtransforms import matplotlib.collections as mcollections -from matplotlib.legend_handler import HandlerTuple +from matplotlib.legend_handler import HandlerTuple, HandlerAnnotation +import matplotlib.patches as mpatches import matplotlib.legend as mlegend @@ -216,6 +217,55 @@ def test_hatching(): ax.legend(handlelength=4, handleheight=4) +@image_comparison(baseline_images=['legend_all_annotation_text'], + extensions=['png'], + style='mpl20') +def test_legend_all_annotation(): + # Related to issue 8236 + # Tests all annotations and text in legend + fig, ax = plt.subplots(1) + ax.plot([0, 1], [0, 0], label='line1') + ax.plot([0, 1], [1, 1], label='line2') + ax.set_xticklabels('') + ax.set_yticklabels('') + # no text, no arrow + ax.annotate("", + xy=(0.1, 0.5), + xytext=(0.1, 0.5), + label='annotation (empty)') + # text, no arrow + my_annotation = ax.annotate("X", + xy=(0.1, 0.5), + xytext=(0.1, 0.5), + color='C2', + label='annotation (text, no arrow)') + # no text, arrow + ax.annotate("", + xy=(0.3, 1.0), + xytext=(0.3, 0.0), + arrowprops={'arrowstyle': '<->', 'color': 'C7'}, + label='annotation (no text, arrow)') + # Fancy arrow patch + arrpatch = mpatches.FancyArrowPatch([0.5, 0.8], [0.9, 0.9], + arrowstyle='<|-', + mutation_scale=20, + color='C3', + label='arrowpatch') + ax.add_patch(arrpatch) + # Long text, will not be used in legend + ax.text(x=0.1, y=0.1, + s='Hello', + color='C5', + label='text') + # Short text, copied in legend + ax.text(x=0.1, y=0.2, + s='Z', + color='C0', + label='short text') + ax.legend(handler_map={my_annotation: HandlerAnnotation(rep_str='Abcde', + rep_maxlen=0)}) + + def test_legend_remove(): fig = plt.figure() ax = fig.add_subplot(1, 1, 1)