matplotlib · blink1073 · Nov 5, 2016 · Nov 5, 2016 · Nov 5, 2016 · Nov 5, 2016
diff --git a/examples/widgets/painter.py b/examples/widgets/painter.py
@@ -0,0 +1,36 @@
+"""
+===============================
+Painter Tool Demo
+===============================
+
+Drag the mouse to paint selected areas of the plot. The callback prints
+the (x, y) coordinates of the center of the painted region.
+
+"""
+from __future__ import print_function
+
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.widgets import Painter
+
+# Create a figure and axes for the plot.
+fig, ax = plt.subplots()
+
+
+# Generate data and create a scatter plot.
+data = np.random.rand(100, 2)
+data[:, 1] *= 2
+pts = ax.scatter(data[:, 0], data[:, 1], s=80)
+
+
+# Define the "on_select" callback.
+def callback(x, y):
+    print("(%3.2f, %3.2f)" % (x, y))
+
+
+print("\n click and drag \n (x, y)")
+
+# Create the painter tool and show the plot.
+p = Painter(ax, callback)
+p.label = 1  # set the color
+plt.show()
diff --git a/lib/matplotlib/tests/test_widgets.py b/lib/matplotlib/tests/test_widgets.py
@@ -271,3 +271,36 @@ def clicked_function():
         pass
     cid = check.on_clicked(clicked_function)
     check.disconnect(cid)
+
+
+@cleanup
+def test_paint_tool():
+    ax = get_ax()
+
+    def onselect(x, y):
+        pass
+
+    tool = widgets.Painter(ax, onselect)
+
+    tool.radius = 10
+    assert tool.radius == 10
+    tool.label = 2
+    assert tool.label == 2
+
+    do_event(tool, 'press', xdata=100, ydata=100)
+    do_event(tool, 'onmove', xdata=110, ydata=110)
+    do_event(tool, 'release')
+
+    # Ensure that the correct number of pixels are labeled.
+    assert tool.overlay[tool.overlay == 2].size == 878
+
+    tool.label = 5
+    do_event(tool, 'press', xdata=20, ydata=20)
+    do_event(tool, 'onmove', xdata=40, ydata=40)
+    do_event(tool, 'release')
+
+    assert tool.overlay[tool.overlay == 5].size == 882
+    assert tool.overlay[tool.overlay == 2].size == 878
+
+    tool.overlay = tool.overlay * 0
+    assert tool.overlay.sum() == 0
diff --git a/lib/matplotlib/widgets.py b/lib/matplotlib/widgets.py
@@ -19,6 +19,7 @@
 import numpy as np
 from matplotlib import rcParams
 
+from .colors import ListedColormap, NoNorm
 from .mlab import dist
 from .patches import Circle, Rectangle, Ellipse
 from .lines import Line2D
@@ -1468,6 +1469,8 @@ def __init__(self, ax, onselect, useblit=False, button=None,
 
         if isinstance(button, int):
             self.validButtons = [button]
+        elif button is None:
+            self.validButtons = [1, 2, 3]
         else:
             self.validButtons = button
 
@@ -2595,3 +2598,197 @@ def onmove(self, event):
             self.canvas.blit(self.ax.bbox)
         else:
             self.canvas.draw_idle()
+
+
+LABELS_CMAP = ListedColormap(['white', 'red', 'dodgerblue', 'gold',
+                              'greenyellow', 'blueviolet'])
+
+
+class Painter(_SelectorWidget):
+    """
+    Paint regions on an axes.
+
+    For the widget to remain responsive you must keep a reference to it.
+
+    This is not meant to be used with polar coordinate plots.
+
+    Resizing the plot will clear the existing selection.
+
+    Example usage::
+
+        import numpy as np
+        import matplotlib.pyplot as plt
+        from matplotlib.widgets import Painter
+
+        data = np.random.rand(100, 2)
+        data[:, 1] *= 2
+
+        fig, ax = plt.subplots()
+        pts = ax.scatter(data[:, 0], data[:, 1], s=80)
+
+        def test(x, y):
+            print(x, y)
+
+        p = Painter(ax, test)
+        plt.show()
+    """
+
+    def __init__(self, ax, on_select=None, overlay_props=None,
+                 cursor_props=None, radius=5, cmap=LABELS_CMAP,
+                 useblit=True, button=None, state_modifier_keys=None):
+        """
+        Parameters
+        ==========
+
+        ax: :class:`~matplotlib.axes.Axes`
+            The parent axes for the widget.
+        on_select: function
+            A callback for when a region is painted. Called with the
+            (x, y) coordinates of the region.
+        overlay_props: dict
+            The properties to apply to the overlay.
+        cursor_props: dcit
+            The properties to apply to the cursor.
+        radius: int
+            The radius of the cursor in pixels.
+        cmap: :class:~matplotlib.colorls.ListedColormap`
+            The colormap to use for the cursors.
+        useblit: bool
+            Whether to use blitting.
+        button: list
+            The button numbers supported for the tool (defaults to [1, 2, 3])
+        state_modifier_keys: dict
+            A mapping of key names to state modifiers.
+        """
+        super(Painter, self).__init__(
+            ax, on_select, useblit=useblit, button=button,
+            state_modifier_keys=None
+        )
+        self.cmap = cmap
+        self._previous = None
+        self._overlay = None
+        self._overlay_plot = None
+        self._cursor_shape = [0, 0, 0]
+
+        props = dict(edgecolor='r', facecolor='0.7', alpha=1,
+                     animated=self.useblit, visible=False, zorder=2)
+        props.update(cursor_props or {})
+        self._cursor = Rectangle((0, 0), 0, 0, **props)
+        self.ax.add_patch(self._cursor)
+
+        x0, x1 = self.ax.get_xlim()
+        y0, y1 = self.ax.get_ylim()
+        if y0 < y1:
+            origin = 'lower'
+        else:
+            origin = 'upper'
+        props = dict(cmap=self.cmap, alpha=0.5, origin=origin,
+                     norm=NoNorm(), visible=False, zorder=1,
+                     extent=(x0, x1, y0, y1), aspect=self.ax.get_aspect())
+        props.update(overlay_props or {})
+
+        self._resize(None)
+        self._overlay_plot = self.ax.imshow(self._overlay, **props)
+        self.artists = [self._cursor, self._overlay_plot]
+        self.connect_event('resize_event', self._resize)
+
+        # These must be called last
+        self.label = 1
+        self.radius = radius
+        self._drawing = True
+        for artist in self.artists:
+            artist.set_visible(True)
+
+    @property
+    def overlay(self):
+        """The paint overlay image"""
+        return self._overlay
+
+    @overlay.setter
+    def overlay(self, image):
+        self._overlay = image
+        if image is None:
+            self.ax.images.remove(self._overlay_plot)
+            self.update()
+            return
+        self._overlay_plot.set_data(image)
+        self._shape = image.shape
+        x0, x1 = self.ax.get_xlim()
+        y0, y1 = self.ax.get_ylim()
+        self._overlay_plot.set_extent((x0, x1, y0, y1))
+        # Update the radii and window.
+        self.radius = self._radius
+        self.update()
+
+    @property
+    def label(self):
+        """The label index"""
+        return self._label
+
+    @label.setter
+    def label(self, value):
+        if value >= self.cmap.N:
+            raise ValueError('Maximum label value = %s' % len(self.cmap - 1))
+        self._label = value
+        self._cursor.set_edgecolor(self.cmap(value))
+
+    @property
+    def radius(self):
+        """The label radius in pixels"""
+        return self._radius
+
+    @radius.setter
+    def radius(self, r):
+        self._radius = r
+        xfm = self.ax.transData.inverted()
+        x0, y0 = xfm.transform((0, 0))
+        x1, y1 = xfm.transform((r, r))
+        self._rx, self._ry = abs(x1 - x0), abs(y1 - y0)
+
+        self._cursor.set_width(self._rx * 2)
+        self._cursor.set_height(self._ry * 2)
+
+    def _press(self, event):
+        self._update_cursor(event.xdata, event.ydata)
+        self._update_overlay(event.x, event.y)
+        self.update()
+
+    def _onmove(self, event):
+        self._update_cursor(event.xdata, event.ydata)
+        if event.button and event.button in self.validButtons:
+            self._update_overlay(event.x, event.y)
+        self.update()
+        self.onselect(event.xdata, event.ydata)
+
+    def _resize(self, event):
+        extents = self.ax.get_window_extent().extents
+        self._offsetx = extents[0]
+        self._offsety = extents[1]
+        self._shape = (int(extents[3] - extents[1]),
+                       int(extents[2] - extents[0]))
+        self._overlay = np.zeros(self._shape, dtype='uint8')
+        if self._overlay_plot:
+            self._overlay_plot.set_data(self._overlay)
+
+    def _release(self, event):
+        pass
+
+    def _update_overlay(self, x, y):
+        col = x - self._offsetx
+        row = y - self._offsety
+
+        h, w = self._shape
+        r = self._radius
+
+        xmin = int(max(0, col - r))
+        xmax = int(min(w, col + r + 1))
+        ymin = int(max(0, row - r))
+        ymax = int(min(h, row + r + 1))
+
+        self._overlay[slice(ymin, ymax), slice(xmin, xmax)] = self.label
+        self._overlay_plot.set_data(self._overlay)
+
+    def _update_cursor(self, x, y):
+        x = x - self._rx
+        y = y - self._ry
+        self._cursor.set_xy((x, y))