matplotlib · SwarnaD · Mar 15, 2018 · Mar 15, 2018 · Mar 15, 2018 · Mar 15, 2018
diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py
@@ -39,6 +39,7 @@
 import matplotlib.colors as colors
 from matplotlib.font_manager import FontProperties
 from matplotlib.lines import Line2D
+from matplotlib.text import Annotation
 from matplotlib.patches import Patch, Rectangle, Shadow, FancyBboxPatch
 from matplotlib.collections import (LineCollection, RegularPolyCollection,
                                     CircleCollection, PathCollection,
@@ -815,7 +816,8 @@ 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(),
+        Annotation: legend_handler.HandlerAnnotation()
         }
 
     # (get|set|update)_default_handler_maps are public interfaces to
@@ -1307,12 +1309,12 @@ def _get_legend_handles(axs, legend_handler_map=None):
     """
     handles_original = []
     for ax in axs:
-        handles_original += (ax.lines + ax.patches +
+        handles_original += (ax.lines + ax.patches + ax.texts +
                              ax.collections + ax.containers)
         # support parasite axes:
         if hasattr(ax, 'parasites'):
             for axx in ax.parasites:
-                handles_original += (axx.lines + axx.patches +
+                handles_original += (axx.lines + axx.patches + ax.texts +
                                      axx.collections + axx.containers)
 
     handler_map = Legend.get_default_handler_map()

diff --git a/lib/matplotlib/legend_handler.py b/lib/matplotlib/legend_handler.py
@@ -33,7 +33,8 @@ 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.text import Text, Annotation
+from matplotlib.patches import Rectangle, FancyArrowPatch
 import matplotlib.collections as mcoll
 import matplotlib.colors as mcolors
 
@@ -728,3 +729,66 @@ 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" which are used by ".Annotations".
+    """
+    def create_artists(self, legend, orig_handle, xdescent, ydescent, width,
+                       height, fontsize, trans):
+        p = Text(x=-xdescent, y=-ydescent,
+                 text=orig_handle.get_text())
+        self.update_prop(p, orig_handle, legend)
+        p.set_transform(trans)
+        p.set_fontsize(0.75*fontsize)
+        c = Rectangle(xy=(-xdescent, -ydescent - (height/5)), width=width, height=7*height/5,
+                      facecolor="none", edgecolor="none")
+        c.set_transform(trans)
+        p.set_clip_path(c)
+        return[p]
+
+
+class HandlerFancyArrowPatch(HandlerBase):
+    """
+    Handler for ".FancyArrow" which are used by ".Annotations".
+    """
+    def update_prop(self, legend_handle, orig_handle, legend):
+        self._update_prop(legend_handle, orig_handle)
+        legend_handle.set_arrowstyle(orig_handle.get_arrowstyle())
+        legend_handle.set_linestyle(orig_handle.get_linestyle())
+        legend_handle.set_mutation_aspect(orig_handle.get_mutation_aspect())
+
+    def create_artists(self, legend, orig_handle, xdescent, ydescent, width,
+                       height, fontsize, trans):
+
+        p = FancyArrowPatch(posA=(-xdescent - (width/10), -ydescent +
+                                  height/2),
+                            posB=(-xdescent + width + (width/10), -ydescent +
+                                  height/2),
+                            mutation_scale=height*1.5)
+        self.update_prop(p, orig_handle, legend)
+        p.set_transform(trans)
+        return [p]
+
+
+class HandlerAnnotation(HandlerBase):
+    """
+    Handler for ".Annotation" instances.
+    """
+    def create_artists(self, legend, orig_handle, xdescent, ydescent, width,
+                       height, fontsize, trans):
+        if (orig_handle.arrow_patch is not None):
+            arrow_handler = HandlerFancyArrowPatch()
+            p = arrow_handler.create_artists(legend, orig_handle.arrow_patch,
+                                             xdescent, ydescent, width, height,
+                                             fontsize, trans)
+        elif (orig_handle.get_text() != ""):
+            text_handler = HandlerText()
+            p = text_handler.create_artists(legend, orig_handle, xdescent,
+                                            ydescent, width, height, fontsize,
+                                            trans)
+        else:
+            p = [Rectangle(xy=(-xdescent, -ydescent), width=width,
+                           height=height, facecolor="none", edgecolor="none")]
+        return p
diff --git a/lib/matplotlib/tests/baseline_images/test_legend/all_arrowstyles.png b/lib/matplotlib/tests/baseline_images/test_legend/all_arrowstyles.png
diff --git a/lib/matplotlib/tests/baseline_images/test_legend/all_linestyles.png b/lib/matplotlib/tests/baseline_images/test_legend/all_linestyles.png
diff --git a/lib/matplotlib/tests/baseline_images/test_legend/annotation_colours.png b/lib/matplotlib/tests/baseline_images/test_legend/annotation_colours.png
diff --git a/lib/matplotlib/tests/baseline_images/test_legend/annotation_no_line.png b/lib/matplotlib/tests/baseline_images/test_legend/annotation_no_line.png
diff --git a/lib/matplotlib/tests/baseline_images/test_legend/annotation_no_line_text.png b/lib/matplotlib/tests/baseline_images/test_legend/annotation_no_line_text.png
diff --git a/lib/matplotlib/tests/baseline_images/test_legend/annotation_text.png b/lib/matplotlib/tests/baseline_images/test_legend/annotation_text.png
diff --git a/lib/matplotlib/tests/baseline_images/test_legend/simple_annotation.png b/lib/matplotlib/tests/baseline_images/test_legend/simple_annotation.png
diff --git a/lib/matplotlib/tests/test_legend.py b/lib/matplotlib/tests/test_legend.py
@@ -522,3 +522,233 @@ def test_legend_title_empty():
     leg = ax.legend()
     assert leg.get_title().get_text() == ""
     assert leg.get_title().get_visible() is False
+
+
+@image_comparison(baseline_images=['simple_annotation'],
+                  extensions=['png'])
+def test_simple_annotation():
+    # tests that a simple example graph with annotations
+    # appear on the legend with their image and label
+    x = np.arange(0.0, 15.0, 0.01)
+    y = np.sin(0.3*np.pi*x)
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+    ax.plot(x, y, lw=2, label="sin(x)")
+    ax.annotate("",
+                xy=(1.6, 1.03), 
+                xytext=(8.4, 1.03), 
+                arrowprops={'arrowstyle':'<->'}, 
+                label="wavelength")
+    ax.annotate("",
+                xy=(8.4, 0), 
+                xytext=(8.35, 1.0), 
+                arrowprops={'arrowstyle':'<->', 'ls':'dashed'}, 
+                label="amplitude")
+    ax.set_ylim(-1.5, 1.5)
+    ax.legend()
+    plt.show()
+
+
+@image_comparison(baseline_images=['all_linestyles'],
+                  extensions=['png'])
+def test_all_linestyles():
+    # tests that annotations with different linestyles
+    # appear on the legend with their image and label
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+    ax.set_xlim(0, 1.7)
+    ax.annotate("",
+                xy=(0.1, 0.1), 
+                xytext=(0.1, 0.9), 
+                arrowprops={'arrowstyle':'-', 'ls':'solid'}, 
+                label="solid")
+    ax.annotate("",
+                xy=(0.3, 0.1), 
+                xytext=(0.3, 0.9), 
+                arrowprops={'arrowstyle':'-', 'ls':'dashed'}, 
+                label="dashed")
+    ax.annotate("",
+                xy=(0.5, 0.1), 
+                xytext=(0.5, 0.9), 
+                arrowprops={'arrowstyle':'-', 'ls':'dashdot'}, 
+                label="dashdot")
+    ax.annotate("",
+                xy=(0.7, 0.1), 
+                xytext=(0.7, 0.9), 
+                arrowprops={'arrowstyle':'-', 'ls':'dotted'}, 
+                label="dotted")    
+    ax.annotate("",
+                xy=(0.9, 0.1), 
+                xytext=(0.9, 0.9), 
+                arrowprops={'arrowstyle':'-', 'ls':(0, (1,5))}, 
+                label="offset, on-off-dash-seq")    
+    ax.legend()
+    plt.show()
+
+
+@image_comparison(baseline_images=['all_arrowstyles'],
+                  extensions=['png'])
+def test_all_arrowstyles():
+    # tests that annotations with different arrowstyles
+    # appear on the legend with their image and label
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+    ax.set_ylim(0, 2.5)
+    ax.annotate("",
+                xy=(0.1, 0.1), 
+                xytext=(0.5, 0.1), 
+                arrowprops={'arrowstyle':'-'}, 
+                label="-")
+    ax.annotate("",
+                xy=(0.1, 0.3), 
+                xytext=(0.5, 0.3), 
+                arrowprops={'arrowstyle':'->'}, 
+                label="->")
+    ax.annotate("",
+                xy=(0.1, 0.5), 
+                xytext=(0.5, 0.5), 
+                arrowprops={'arrowstyle':'-['}, 
+                label="-[") 
+    ax.annotate("",
+                xy=(0.1, 0.7), 
+                xytext=(0.5, 0.7), 
+                arrowprops={'arrowstyle':'|-|'}, 
+                label="|-|")
+    ax.annotate("",
+                xy=(0.1, 0.9), 
+                xytext=(0.5, 0.9), 
+                arrowprops={'arrowstyle':'-|>'}, 
+                label="-|>")
+    ax.annotate("",
+                xy=(0.1, 1.1), 
+                xytext=(0.5, 1.1), 
+                arrowprops={'arrowstyle':'<-'}, 
+                label="<-")
+    ax.annotate("",
+                xy=(0.1, 1.3), 
+                xytext=(0.5, 1.3), 
+                arrowprops={'arrowstyle':'<->'}, 
+                label="<->")
+    ax.annotate("",
+                xy=(0.1, 1.5), 
+                xytext=(0.5, 1.5), 
+                arrowprops={'arrowstyle':'<|-'}, 
+                label="<|-")
+    ax.annotate("",
+                xy=(0.1, 1.7), 
+                xytext=(0.5, 1.7), 
+                arrowprops={'arrowstyle':'<|-|>'}, 
+                label="<|-|>")
+    ax.annotate("",
+                xy=(0.1, 1.9), 
+                xytext=(0.5, 1.9), 
+                arrowprops={'arrowstyle':'fancy'}, 
+                label="fancy")
+    ax.annotate("",
+                xy=(0.1, 2.1), 
+                xytext=(0.5, 2.1), 
+                arrowprops={'arrowstyle':'simple'}, 
+                label="simple")
+    ax.annotate("",
+                xy=(0.1, 2.3), 
+                xytext=(0.5, 2.3), 
+                arrowprops={'arrowstyle':'wedge'}, 
+                label="wedge")    
+    ax.legend()
+    plt.show()
+
+
+@image_comparison(baseline_images=['annotation_colours'],
+                  extensions=['png'])
+def test_annotation_colours():
+    # tests that annotations with different colors
+    # appear on the legend with their image and label
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+    ax.annotate("",
+                xy=(0.1, 0.1), 
+                xytext=(0.1, 0.9), 
+                arrowprops={'arrowstyle':'<|-|>', 
+                            'ls':'solid', 
+                            'color':'blue'}, 
+                label="blue")
+    ax.annotate("",
+                xy=(0.2, 0.1), 
+                xytext=(0.2, 0.9), 
+                arrowprops={'arrowstyle':'-',
+                            'ls':'dashed',
+                            'color':'red'}, 
+                label="red")
+    ax.annotate("",
+                xy=(0.3, 0.1), 
+                xytext=(0.3, 0.9), 
+                arrowprops={'arrowstyle':'|-|',
+                            'ls':'dashdot',
+                            'color':'yellow'},
+                label="yellow")    
+    ax.legend()
+    plt.show()
+
+
+@image_comparison(baseline_images=['annotation_text'],
+                  extensions=['png'])
+def test_annotation_text():
+    # tests that annotations with different texts
+    # appear on the legend with their image and label
+    # note: the text itself does not appear on the legend,
+    # only the arrow. The text itself is self-explanatory.
+    # However, if only text is passed, the image will be
+    # the text.
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+    ax.annotate("hello",
+                xy=(0.1, 0.9), 
+                xytext=(0.1, 0.1), 
+                arrowprops={'arrowstyle':'-',
+                            'ls':'dashed',
+                            'color':'green'},
+                label="hello") 
+    ax.annotate("world",
+                xy=(0.3, 0.9), 
+                xytext=(0.3, 0.1), 
+                arrowprops={'arrowstyle':'<->',
+                            'ls':'dashdot',
+                            'color':'purple'},
+                label="world")
+    ax.annotate("short text",
+                xy=(0.5, 0.9), 
+                xytext=(0.5, 0.1),
+                label="short text")    
+    ax.annotate("long text",
+                xy=(0.7, 0.9), 
+                xytext=(0.7, 0.1),
+                label="long text")    
+
+    ax.legend() 
+    plt.show()
+
+
+@image_comparison(baseline_images=['annotation_no_line_text'],
+                  extensions=['png'])
+def test_annotation_no_line_text():
+    # tests that annotations with no line, text, or both
+    # appear on the legend with their image and label
+    # note: if no text, it will not appear in the legend
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+    ax.annotate("",
+                xy=(0.1, 0.9), 
+                xytext=(0.1, 0.1), 
+                arrowprops={'arrowstyle':'-',
+                            'ls':'dashed'},
+                label="no text") 
+    ax.annotate("no line",
+                xy=(0.3, 0.1), 
+                xytext=(0.3, 0.9),
+                label="no line")
+    ax.annotate("",
+                xy=(0.5, 0.1), 
+                xytext=(0.5, 0.9),
+                label="no line or text")
+    ax.legend() 
+    plt.show()