simon effect (psychology experiment ) : build the function of wait4key()

##
#the real experiment for simon effect
#load the library which is our need
import pygame
import sys,random
from pygame.locals import *

pygame.init()
win = pygame.display.set_mode((800,600),DOUBLEBUF|HWSURFACE)
left = (200,300)
right = (600,300)
red  =(255,0,0)
blue = (0,0,255)
black = (0,0,0)

#wait for the pressed key

def wait4key(duration,key_list):
    """the function is to wait for subject to press the right key , and the experient will continue untill thesubject press the key
duration : the time which is the subject wait to press the key ,if the time is too long ,the experiment willcontinue automaticly.
key_list: the key which the subject need to press to continue the experiment, and the key_list must be a list such as [K_A,K_/]"""
    
    fake_bool =False
    time_out = False

    #empty the event before 
    pygame.event.clear()
    
    #use for count the time
    start_time = pygame.time.get_ticks()
    
    #if the subject does not press the right key, he will in the while loop all the time untill the duration is exhausted
    while  not (fake_bool or time_out):
        end_time = pygame.time.get_ticks()
        #if the duration is too long , the experiment will continue
        if end_time - start_time > duration:
            time_out = True
        #if subject press the right key ,the experimet will continue
        for i in pygame.event.get():
            if i.type ==KEYDOWN:
                if i.key in key_list:
                    #prepare for the result of the function
                    response_time = pygame.time.get_ticks()
                    key_name = pygame.key.name(i.key)
                    #if has the right key ,the loop will quit
                    fake_bool = True
    #in the end, if subject press the key ,we will collect the time , and the name of the key
    if fake_bool:
        return start_time, response_time, key_name
    #the purpose of the next line is stay the same with the result
    else:
        return start_time, None, None
        
#experiment for 10 times
for i in range(10):
    r1 = random.choice([0,1])
    if r1 == 0:
        pos = left
    else:
        pos = right

    r2 = random.choice([0,1])
    if r2 == 0:
        color = red
    else:
        color = blue
        
    win.fill(black)
    pygame.draw.circle(win,color,pos,20,0)
    pygame.display.flip()
    result = wait4key(2000,[K_z,K_SLASH])
    print(result)
    #when subject press the key , we will black the window for 0.5 seconds until next stimuli appears
    win.fill(black)
    pygame.display.flip()
    pygame.time.delay(500)