Two seperate key strokes to save and quit a while loop

Question

I have a while loop which for each iteration acquires a new data set. In the following example this is simulated using a random number generator. My while loop is broken by pressing the enter key. This is achieved by the heardEnter() function I found here. This works perfectly.

However, what I also would like to be able to do is save my data by pressing a key in the case that I do not want to exit the loop. The key that I have in mind is either the "s"-key or the spacebar. I have tried to use raw_input() for this. However, the required "Enter" to finalize the raw input makes the combination with the heardEnter function troublesome.

A minimal working example:

import sys
import numpy
import select
import matplotlib.pyplot as plt

def heardEnter():
  i, o, e = select.select([sys.stdin], [], [], 0.0001)
  for s in i:
    if s == sys.stdin:
      input = sys.stdin.readline()
      return True
  return False

x = numpy.linspace(0, 8*numpy.pi, 1000)
y = numpy.cos(x)

plt.ion()
plt.figure()
plt.plot(x, y)
plt.draw()

cont = True
while cont:
  noise = numpy.random.normal(0, 1, len(y))
  y_new = y + noise

  plt.cla()

  plt.plot(x, y_new)

  plt.draw()
  plt.pause(1.E-6)

  if heardEnter():
    cont = False
    plt.close()

Answer 1

You should modify your heardEnter function to work with more than just Enter:

SAVE = 'save'
ENTER = 'enter'

def heardEnter():
  i, o, e = select.select([sys.stdin], [], [], 0.0001)
  for s in i:
    if s == sys.stdin:
      input = sys.stdin.readline()
      if 's' in input:
         return SAVE
      else:
         return ENTER
  return None

Then in your while loop:

  result = heard_enter()
  if result == ENTER:
    cont = False
    plt.close()
  elif result == SAVE:
    # do save

Since heardEnter needs to communicate more info, we can't have it return a simple boolean anymore, and work with constants as a return value instead. This will allow you to work with additional keys in the future as well.