Legend refactoring. Drawing of legends is done by LegendHandlers

leejjoon · leejjoon · commit 572ef7e27a49 · 2011-03-02T21:09:39.000+09:00
diff --git a/examples/pylab_examples/legend_demo_custom_handler.py b/examples/pylab_examples/legend_demo_custom_handler.py
@@ -0,0 +1,30 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+ax = plt.subplot(111, aspect=1)
+x, y = np.random.randn(2, 20)
+#[1.1, 2, 2.8], [1.1, 2, 1.8]
+l1, = ax.plot(x,y, "k+", mew=3, ms=12)
+l2, = ax.plot(x,y, "w+", mew=1, ms=10)
+
+import matplotlib.patches as mpatches
+c = mpatches.Circle((0, 0), 1, fc="g", ec="r", lw=3)
+ax.add_patch(c)
+
+
+
+from matplotlib.legend_handler import HandlerPatch, HandlerMulti
+
+def make_legend_ellipse(legend, orig_handle,
+                        xdescent, ydescent,
+                        width, height, fontsize):
+    p = mpatches.Ellipse(xy=(0.5*width-0.5*xdescent, 0.5*height-0.5*ydescent),
+                         width = width+xdescent, height=(height+ydescent))
+
+    return p
+
+plt.legend([c, "handle 2"], ["Label 1", "Label 2"],
+           handler_map={mpatches.Circle:HandlerPatch(patch_func=make_legend_ellipse),
+                        "handle 2":HandlerMulti(l1, l2)})
+
+plt.show()
diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py
@@ -37,6 +37,8 @@
 
 from matplotlib.offsetbox import HPacker, VPacker, TextArea, DrawingArea, DraggableOffsetBox
 
+import legend_handler
+
 
 class DraggableLegend(DraggableOffsetBox):
     def __init__(self, legend, use_blit=False, update="loc"):
@@ -167,6 +169,7 @@ def __init__(self, parent, handles, labels,
                  bbox_to_anchor = None, # bbox that the legend will be anchored.
                  bbox_transform = None, # transform for the bbox
                  frameon = True, # draw frame
+                 handler_map = None,
                  ):
         """
         - *parent* : the artist that contains the legend
@@ -236,6 +239,9 @@ def __init__(self, parent, handles, labels,
         self.legendHandles = []
         self._legend_title_box = None
 
+
+        self._handler_map = handler_map
+
         localdict = locals()
 
         for name in propnames:
@@ -470,6 +476,36 @@ def _approx_text_height(self, renderer=None):
             return renderer.points_to_pixels(self._fontsize)
 
 
+    _default_handler_map = {Line2D:legend_handler.HandlerLine2D(),
+                            Patch:legend_handler.HandlerPatch(),
+                            LineCollection:legend_handler.HandlerLineCollection(),
+                            RegularPolyCollection:legend_handler.HandlerRegularPolyCollection(),
+                            CircleCollection:legend_handler.HandlerCircleCollection()
+                            }
+
+    def get_legend_handler_map(self):
+        if self._handler_map:
+            hm = self._default_handler_map.copy()
+            hm.update(self._handler_map)
+            return hm
+        else:
+            return self._default_handler_map
+
+    def get_legend_handler(self, legend_handler_map, orig_handle):
+            if orig_handle in legend_handler_map:
+                handler = legend_handler_map[orig_handle]
+            else:
+                
+                for handle_type in type(orig_handle).mro():
+                    if handle_type in legend_handler_map:
+                        handler = legend_handler_map[handle_type]
+                        break
+                else:
+                    handler = None
+
+            return handler
+
+        
     def _init_legend_box(self, handles, labels):
         """
         Initiallize the legend_box. The legend_box is an instance of
@@ -516,154 +552,37 @@ def _init_legend_box(self, handles, labels):
         # manually set their transform to the self.get_transform().
 
 
-        for handle, lab in zip(handles, labels):
-            if isinstance(handle, RegularPolyCollection) or \
-                   isinstance(handle, CircleCollection):
-                npoints = self.scatterpoints
-            else:
-                npoints = self.numpoints
-            if npoints > 1:
-                # we put some pad here to compensate the size of the
-                # marker
-                xdata = np.linspace(0.3*fontsize,
-                                    (self.handlelength-0.3)*fontsize,
-                                    npoints)
-                xdata_marker = xdata
-            elif npoints == 1:
-                xdata = np.linspace(0, self.handlelength*fontsize, 2)
-                xdata_marker = [0.5*self.handlelength*fontsize]
-
-            if isinstance(handle, Line2D):
-                ydata = ((height-descent)/2.)*np.ones(xdata.shape, float)
-                legline = Line2D(xdata, ydata)
-
-                legline.update_from(handle)
-                self._set_artist_props(legline) # after update
-                legline.set_clip_box(None)
-                legline.set_clip_path(None)
-                legline.set_drawstyle('default')
-                legline.set_marker('None')
-
-                handle_list.append(legline)
-
-                legline_marker = Line2D(xdata_marker, ydata[:len(xdata_marker)])
-                legline_marker.update_from(handle)
-                self._set_artist_props(legline_marker)
-                legline_marker.set_clip_box(None)
-                legline_marker.set_clip_path(None)
-                legline_marker.set_linestyle('None')
-                if self.markerscale !=1:
-                    newsz = legline_marker.get_markersize()*self.markerscale
-                    legline_marker.set_markersize(newsz)
-                # we don't want to add this to the return list because
-                # the texts and handles are assumed to be in one-to-one
-                # correpondence.
-                legline._legmarker = legline_marker
-
-            elif isinstance(handle, Patch):
-                p = Rectangle(xy=(0., 0.),
-                              width = self.handlelength*fontsize,
-                              height=(height-descent),
-                              )
-                p.update_from(handle)
-                self._set_artist_props(p)
-                p.set_clip_box(None)
-                p.set_clip_path(None)
-                handle_list.append(p)
-            elif isinstance(handle, LineCollection):
-                ydata = ((height-descent)/2.)*np.ones(xdata.shape, float)
-                legline = Line2D(xdata, ydata)
-                self._set_artist_props(legline)
-                legline.set_clip_box(None)
-                legline.set_clip_path(None)
-                lw = handle.get_linewidth()[0]
-                dashes = handle.get_dashes()[0]
-                color = handle.get_colors()[0]
-                legline.set_color(color)
-                legline.set_linewidth(lw)
-                if dashes[0] is not None: # dashed line
-                    legline.set_dashes(dashes[1])
-                handle_list.append(legline)
-
-            elif isinstance(handle, RegularPolyCollection):
-
-                #ydata = self._scatteryoffsets
-                ydata = height*self._scatteryoffsets
-
-                size_max, size_min = max(handle.get_sizes())*self.markerscale**2,\
-                                     min(handle.get_sizes())*self.markerscale**2
-                if self.scatterpoints < 4:
-                    sizes = [.5*(size_max+size_min), size_max,
-                             size_min]
-                else:
-                    sizes = (size_max-size_min)*np.linspace(0,1,self.scatterpoints)+size_min
-
-                p = type(handle)(handle.get_numsides(),
-                                 rotation=handle.get_rotation(),
-                                 sizes=sizes,
-                                 offsets=zip(xdata_marker,ydata),
-                                 transOffset=self.get_transform(),
-                                 )
-
-                p.update_from(handle)
-                p.set_figure(self.figure)
-                p.set_clip_box(None)
-                p.set_clip_path(None)
-                handle_list.append(p)
-
-            elif isinstance(handle, CircleCollection):
-
-                ydata = height*self._scatteryoffsets
-
-                size_max, size_min = max(handle.get_sizes())*self.markerscale**2,\
-                                     min(handle.get_sizes())*self.markerscale**2
-                if self.scatterpoints < 4:
-                    sizes = [.5*(size_max+size_min), size_max,
-                             size_min]
-                else:
-                    sizes = (size_max-size_min)*np.linspace(0,1,self.scatterpoints)+size_min
-
-                p = type(handle)(sizes,
-                                 offsets=zip(xdata_marker,ydata),
-                                 transOffset=self.get_transform(),
-                                 )
-
-                p.update_from(handle)
-                p.set_figure(self.figure)
-                p.set_clip_box(None)
-                p.set_clip_path(None)
-                handle_list.append(p)
-            else:
-                handle_type = type(handle)
-                warnings.warn("Legend does not support %s\nUse proxy artist instead.\n\nhttp://matplotlib.sourceforge.net/users/legend_guide.html#using-proxy-artist\n" % (str(handle_type),))
-                handle_list.append(None)
+        legend_handler_map = self.get_legend_handler_map()
 
+        for orig_handle, lab in zip(handles, labels):
 
+            handler = self.get_legend_handler(legend_handler_map, orig_handle)
 
-            handle = handle_list[-1]
-            if handle is not None: # handle is None is the artist is not supproted
-                textbox = TextArea(lab, textprops=label_prop,
-                                   multilinebaseline=True, minimumdescent=True)
-                text_list.append(textbox._text)
+            if handler is None:
+                warnings.warn("Legend does not support %s\nUse proxy artist instead.\n\nhttp://matplotlib.sourceforge.net/users/legend_guide.html#using-proxy-artist\n" % (str(orig_handle),))
+                handle_list.append(None)
+                continue
+            
+            textbox = TextArea(lab, textprops=label_prop,
+                               multilinebaseline=True, minimumdescent=True)
+            text_list.append(textbox._text)
 
-                labelboxes.append(textbox)
+            labelboxes.append(textbox)
 
-                handlebox = DrawingArea(width=self.handlelength*fontsize,
-                                        height=height,
-                                        xdescent=0., ydescent=descent)
+            handlebox = DrawingArea(width=self.handlelength*fontsize,
+                                    height=height,
+                                    xdescent=0., ydescent=descent)
 
-                handlebox.add_artist(handle)
+            handle = handler(self, orig_handle, \
+                             #xdescent, ydescent, width, height,
+                             fontsize,
+                             handlebox)
+            handle_list.append(handle)
+                
 
-                # special case for collection instances
-                if isinstance(handle, RegularPolyCollection) or \
-                       isinstance(handle, CircleCollection):
-                    handle._transOffset = handlebox.get_transform()
-                    handle.set_transform(None)
 
 
-                if hasattr(handle, "_legmarker"):
-                    handlebox.add_artist(handle._legmarker)
-                handleboxes.append(handlebox)
+            handleboxes.append(handlebox)
 
 
         if len(handleboxes) > 0:
diff --git a/lib/matplotlib/legend_handler.py b/lib/matplotlib/legend_handler.py
