diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py index dc7b82386216..65f32ecb4c4e 100644 --- a/lib/matplotlib/axes/_axes.py +++ b/lib/matplotlib/axes/_axes.py @@ -263,7 +263,8 @@ def _get_legend_handles(self, legend_handler_map=None): """ handles_original = (self.lines + self.patches + - self.collections + self.containers) + self.collections + self.containers + + self.texts) handler_map = mlegend.Legend.get_default_handler_map() if legend_handler_map is not None: diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py index 5165db5c824d..92b66ac0aff7 100644 --- a/lib/matplotlib/legend.py +++ b/lib/matplotlib/legend.py @@ -36,7 +36,8 @@ from matplotlib.cbook import is_string_like, silent_list, is_hashable 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.text import Text, Annotation from matplotlib.collections import (LineCollection, RegularPolyCollection, CircleCollection, PathCollection, PolyCollection) @@ -493,7 +494,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(), + Text:legend_handler.HandlerText(), + FancyArrowPatch:legend_handler.HandlerFancyArrowPatch(), + Annotation:legend_handler.HandlerAnnotation(), } # (get|set|update)_default_handler_maps are public interfaces to diff --git a/lib/matplotlib/legend_handler.py b/lib/matplotlib/legend_handler.py index 67706de0486c..3959340a3802 100644 --- a/lib/matplotlib/legend_handler.py +++ b/lib/matplotlib/legend_handler.py @@ -34,11 +34,13 @@ 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 +from matplotlib.text import Text, Annotation import matplotlib.collections as mcoll import matplotlib.colors as mcolors + def update_from_first_child(tgt, src): tgt.update_from(src.get_children()[0]) @@ -256,6 +258,21 @@ def create_artists(self, legend, orig_handle, 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 HandlerLineCollection(HandlerLine2D): """ Handler for LineCollection instances. @@ -578,18 +595,30 @@ class HandlerTuple(HandlerBase): The number of sections to divide the legend area into. If None, use the length of the input tuple. Default is 1. - 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 + The relative width of sections. 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 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 - self._ndivide = ndivide - self._pad = pad HandlerBase.__init__(self, **kwargs) def create_artists(self, legend, orig_handle, @@ -608,19 +637,30 @@ 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 = [xdescent - (width + pad) * i for i in range(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) + 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, six.next(xds_cycle), ydescent, - width, height, + six.next(widths_cycle), + height, fontsize, trans) a_list.extend(_a_list) @@ -661,3 +701,109 @@ def create_artists(self, legend, orig_handle, self.update_prop(p, orig_handle, legend) p.set_transform(trans) return [p] + + +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 text and arrow sections. 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)