How to pause a pylab figure until a key is pressed or mouse is clicked?

Question

I am trying to program an animated simulation using pylab and networkx. The simulation is not interesting all the time so most of the time I want it to go fast, however, I want to be able to pause it and look at it when it looks interesting. Pausing the screen until keypress will solve my problem, because I can press the key as fast/slow as I want.

Here's an example situation:

import numpy as np
import networkx as nx
import pylab as plt
import sys

def drawGraph(matrix):
    plt.clf()
    G = nx.DiGraph(np.array(matrix))
    nx.draw_networkx(G)
    plt.draw()

    plt.pause(1) #I want this pause to be replaced by a keypress
    #so that it pauses as long as I want

A=[[0,1],[1,0]]
B=[[0,1],[0,0]]
x=1
while True:
    if x==1:
        drawGraph(A)
        x=0
    else:
        drawGraph(B)
        x=1

How should I rewrite the plt.pause(1) line, so that the program pauses until keypress?

Some approaches suggested in other threads pauses the program, but the picture disappears or doesn't update.

Answer 1

Is there some reason not to use waitforbuttonpress()?

import matplotlib.pyplot as plt

A=[[0,1],[1,0]]
B=[[0,1],[0,0]]
x=1
while True:
    if x==1:
        drawGraph(A)
        x=0
    else:
        drawGraph(B)
        x=1
    plt.waitforbuttonpress()

This will, as it says, wait for a key or button press for something to happen. It returns values if you want to know more about the event. Very easy.

Answer 2

The following code stops and restarts with a mouse click. It uses the "TKAgg" backend:

import numpy as np
import networkx as nx
import matplotlib
matplotlib.use("TkAgg")

import pylab as plt
plt.ion()

fig = plt.figure()

pause = False

def onclick(event):
    global pause
    pause = not pause

fig.canvas.mpl_connect('button_press_event', onclick)

def drawGraph(matrix):
    fig.clear()
    G = nx.DiGraph(np.array(matrix))
    nx.draw_networkx(G)
    plt.draw()

A=[[0,1],[1,0]]
B=[[0,1],[0,0]]
x=1
while True:
    if not pause:
        if x==1:
            drawGraph(A)
            x=0
        else:
            drawGraph(B)
            x=1
    fig.canvas.get_tk_widget().update() # process events