Help Needed in Countdown Function within Larger Trial Function

I’m coding an fMRI task that requires a response from participants within a 5 second window. This response is nestled within a larger function (pos_effort_trial_1) that represents the entire trial: a participant is shown a stimulus, told to respond, and then when prompted they have 5 seconds to respond. Based off a previous post’s suggestions, I created a countdown timer with CODE 1 and created the function “text_squeeze”. However, I wanted to use psychopy.event.waitKeys() to get a response from the participant during these 5 seconds. This key press (named “response”) then populates an if/else statement for pos_effort_trial_1. This is where the problem occurs. Because the “text_squeeze” function doesn’t allow the clock to countdown unless win.flip() is embedded in the definition of the function, “response” has to also be a part of the function. When it was not included in the function, “response” couldn’t collect a button press until after text_squeeze timed out. However, when “response” was within “text_squeeze,” the value within “response” didn’t carry through to the if/else statement in pos_effort_trial_1. I tried putting the rest of the pos_effort_trial_1 if/else statement within the “text_squeeze” function, but when I added this additional code, the countdown clock again stopped progressing and there were other glitches in my code. Is there a way to keep “response” within “text_squeeze” but have the rest of the pos_effort_trial_1 work and use “response”? Code is below.

''''CODE 1: ''''
#define window 
win = psychopy.visual.Window([1920, 1080], fullscr = True,  
    monitor = 'testMonitor', color = [-1, -1, -1], units="pix") 
kb = keyboard.Keyboard() 

clock = core.Clock() 
#countdown clock 

def text_squeeze(win_in, text_in, color_in): 
    '''function that shows squeeze and countdown''' 
    message = visual.TextStim(win_in, text = '', color = color_in, height=20, pos=(0,150)) 
    clock.reset() 
    while clock.getTime()<5: 
        message.text = 'Time left: \n' + str(5-clock.getTime()) 
        message.draw() 
        win.flip() 

#calling in the function
text_squeeze(win, 'text', 'white') 


''''CODE 2: with response in a way that DOES collect a response during countdown''' 
def text_squeeze(win_in, text_in, color_in): 
    '''function that shows squeeze and countdown''' 
    message = visual.TextStim(win_in, text = '', color = color_in, height=20, pos=(0,150)) 
    clock.reset() 
    while clock.getTime()<5: 
        message.text = 'Time left: \n' + str(5-clock.getTime()) 
        message.draw() 
        win.flip() 
        response = event.waitKeys(maxWait = 4.5, keyList = ['5'])

'''CODE 3: with response in a way that DOES NOT collect a response during countdown'''

#calling in the function
text_squeeze(win, 'text', 'white') 
response = event.waitKeys(maxWait = 4.5, keyList = ['5'])

'''CODE 4: if/else intro statements following text_squeeze
if response is None: 
 [...]
else: 
 [...]