Add legend handler for Annotation

Wojciech-Lange · oscargus · commit e4ee185305b5 · 2022-05-29T17:56:22.000+02:00
diff --git a/examples/text_labels_and_annotations/legend.py b/examples/text_labels_and_annotations/legend.py
@@ -21,7 +21,11 @@
 ax.plot(a, d, 'k:', label='Data length')
 ax.plot(a, c + d, 'k', label='Total message length')
 
-legend = ax.legend(loc='upper center', shadow=True, fontsize='x-large')
+#Create an arrow with pre-defined label.
+ax.annotate("", xy=(1.5, 4.5), xytext=(1.5, 9.0),
+            arrowprops={'arrowstyle':'<->', 'color':'C7'}, label='Distance')
+
+legend = ax.legend(loc='upper center', shadow=True, fontsize='large')
 
 # Put a nicer background color on the legend.
 legend.get_frame().set_facecolor('C0')
diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py
@@ -34,11 +34,11 @@
 from matplotlib.font_manager import FontProperties
 from matplotlib.lines import Line2D
 from matplotlib.patches import (Patch, Rectangle, Shadow, FancyBboxPatch,
-                                StepPatch)
+                                StepPatch, FancyArrowPatch)
 from matplotlib.collections import (
     Collection, CircleCollection, LineCollection, PathCollection,
     PolyCollection, RegularPolyCollection)
-from matplotlib.text import Text
+from matplotlib.text import Annotation, Text
 from matplotlib.transforms import Bbox, BboxBase, TransformedBbox
 from matplotlib.transforms import BboxTransformTo, BboxTransformFrom
 from matplotlib.offsetbox import (
@@ -649,7 +649,9 @@ def draw(self, renderer):
             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(),
+        Annotation: legend_handler.HandlerAnnotation()
         }
 
     # (get|set|update)_default_handler_maps are public interfaces to
@@ -799,6 +801,7 @@ def _init_legend_box(self, handles, labels, markerfirst=True):
                                              self._legend_handle_box])
         self._legend_box.set_figure(self.figure)
         self._legend_box.axes = self.axes
+        self._legend_box.set_offset(self._findoffset)
         self.texts = text_list
         self.legendHandles = handle_list
 
diff --git a/lib/matplotlib/legend_handler.py b/lib/matplotlib/legend_handler.py
@@ -34,7 +34,7 @@ def legend_artist(self, legend, orig_handle, fontsize, handlebox)
 
 from matplotlib import _api, cbook
 from matplotlib.lines import Line2D
-from matplotlib.patches import Rectangle
+from matplotlib.patches import Rectangle, FancyArrowPatch
 import matplotlib.collections as mcoll
 
 
@@ -806,3 +806,54 @@ 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 HandlerAnnotation(HandlerBase):
+    """
+    Handler for `.Annotation` instances.
+    Defers to `HandlerFancyArrowPatch` to draw the annotation arrow (if any).
+
+    Parameters
+    ----------
+    pad : float, optional
+        If None, fall back to :rc:`legend.borderpad` as the default.
+        In units of fraction of font size. Default is None
+        .
+    width_ratios : two-tuple, optional
+        The relative width of the respective text/arrow legend annotation pair.
+        Default is (1, 4).
+    """
+
+    def __init__(self, pad=None, width_ratios=(1, 4), **kwargs):
+
+        self._pad = pad
+        self._width_ratios = width_ratios
+
+        HandlerBase.__init__(self, **kwargs)
+
+    def create_artists(self, legend, orig_handle, xdescent, ydescent, width,
+                       height, fontsize, trans):
+
+        if orig_handle.arrow_patch is not None:
+
+            # Arrow without text
+
+            handler = HandlerFancyArrowPatch()
+            handle = orig_handle.arrow_patch
+
+        return handler.create_artists(legend, handle, xdescent, ydescent,
+                                      width, height, fontsize, trans)
