How to make class for onclick(event) functions for Interactive plot in matplotlib?

Question

I have tried to create an interactive matplotlib plot using some functions. I want to group the functions into one class (I am still new to this, took help from someone else's code)

import matplotlib.pyplot as plt


def draw_line(startx,starty):

    ax = plt.gca()
    xy = plt.ginput(1)
    x = [startx,xy[0][0]]
    y = [starty,xy[0][1]]
    line = ax.plot(x,y, picker=True , pickradius = 5 , color = "blue")
    ax.figure.canvas.draw()        


def onclick(event):

    """
    This implements click functionality.  If it's a double click do something,
    else ignore.
    Once in the double click block, if its a left click, wait for a further 
    click and draw a line between the double click co-ordinates and that click
    (using ginput(1) - the 1 means wait for one mouse input - a higher number
    is used to get multiple clicks to define a polyline)
    """
    ax = plt.gca()

    if event.dblclick:
        if event.button == 1:
            # Draw line
            draw_line(event.xdata,event.ydata) # here you click on the plot
        else:
            pass # Do nothing

    if event.button == 1:
        pass


def onpick(event):    

    ax = plt.gca()

    """

    Handles the pick event - if an object has been picked, store a
    reference to it.  We do this by simply adding a reference to it
    named 'stored_pick' to the axes object.  Note that in python we
    can dynamically add an attribute variable (stored_pick) to an 
    existing object - even one that is produced by a library as in this
    case

    """

    this_artist = event.artist  # the picked object is available as event.artist
    ax.picked_object = this_artist

def on_key(event):

    """
    Function to be bound to the key press event
    If the key pressed is delete and there is a picked object,
    remove that object from the canvas
    """

    if event.key == u'delete':
        ax = plt.gca()
        if ax.picked_object:
            ax.picked_object.remove()
            ax.picked_object = None
            ax.figure.canvas.draw()


def applyplt():

    fig = plt.gcf()
    ax = plt.gca()

    cidonclic = fig.canvas.mpl_connect('button_press_event', onclick)
    cidonpic = fig.canvas.mpl_connect('pick_event', onpick)
    cidonkey = fig.canvas.mpl_connect('key_press_event', on_key)

"""
Basic Plot to test the function.
"""

fig1 = plt.figure(figsize = (10,10))
gs = fig1.add_gridspec(10,10)
ax101 = fig1.add_subplot(gs[:,:])
ax101.set_ylim(0,10)
ax101.set_xlim(0,10)

applyplt()

plt.show()

I want to group these event functions in one class name(object) (e.g.: class Drawer(object))

If any other optimization can be done, please suggest that too. Thanks!

Answer 1

Here are the functions grouped into a class:

import matplotlib.pyplot as plt


class Interactivity:
    def __init__(self, fig = None):
        self.fig = plt.gcf() if fig is None else fig
        self.ax = self.fig.gca()
        self.connections = ()

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.disconnect()

    def connect(self):
        """ Install the event handlers for the plot. """
        self.connections = (
            self.fig.canvas.mpl_connect('button_press_event', self.onclick),
            self.fig.canvas.mpl_connect('pick_event', self.onpick),
            self.fig.canvas.mpl_connect('key_press_event', self.on_key),
        )

    def disconnect(self):
        """ Uninstall the event handlers for the plot. """
        for connection in self.connections:
            self.fig.canvas.mpl_disconnect(connection)

    def draw_line(self, startx, starty):
        xy = plt.ginput(1)
        x = [startx, xy[0][0]]
        y = [starty, xy[0][1]]
        self.ax.plot(x, y, picker=True, pickradius=5, color='blue')
        self.ax.figure.canvas.draw()

    def onclick(self, event):
        """
        This implements click functionality. If it's a double click do
        something, else ignore.
        Once in the double click block, if its a left click, wait for a further
        click and draw a line between the double click co-ordinates and that
        click (using ginput(1) - the 1 means wait for one mouse input - a
        higher number is used to get multiple clicks to define a polyline)
        """
        print('onclick')
        if event.dblclick:
            if event.button == 1:
                self.draw_line(event.xdata, event.ydata)


    def onpick(self, event):
        """
        Handles the pick event - if an object has been picked, store a
        reference to it.  We do this by simply adding a reference to it
        named 'picked_object' to the axes object.
        """
        print('onpick')
        this_artist = event.artist
        # the picked object is available as event.artist
        self.ax.picked_object = this_artist

    def on_key(self, event):
        """
        Function to be bound to the key press event
        If the key pressed is delete and there is a picked object,
        remove that object from the canvas
        """
        print('onkey: ', event.key)
        if event.key == 'delete' and self.ax.picked_object:
            self.ax.picked_object.remove()
            self.ax.picked_object = None
            self.ax.figure.canvas.draw()

Usage:

# Basic plot to test the functionality
fig = plt.figure(figsize = (10,10))
gs = fig.add_gridspec(10,10)
ax101 = fig.add_subplot(gs[:,:])
ax101.set_ylim(0,10)
ax101.set_xlim(0,10)
with Interactivity():
    plt.show()

As you can see, it can be used as a context handler to install and uninstall the handlers automatically, or you can create an instance of it, and then call its connect method to install the handlers manually (and then later optionally call the disconnect method to uninstall the handlers).