Pygame: Button run function that returns a variable

Question

I copied simple button function in pygame module:

def button(text, x, y, w, h, ia_c, a_c, action=None):
    mouse = pygame.mouse.get_pos()
    click = pygame.mouse.get_pressed()
    if x+w > mouse[0] > x and y+h > mouse[1] > y:
        pygame.draw.rect(screen, a_c, (x, y, w, h))
        if click[0] == 1 and action is not None:
            action()
    else:
        pygame.draw.rect(screen, ia_c, (x, y, w, h))
    message_display(text, font_small, (x+(w//2)) * 2, (y+(h//2)) * 2, white)

Then I wanted to create a button, that will run a function, which return a pygame.Surface (image) and assign it to variable:

button("START!", 350, 350, 300, 50, blue, blue_light, which_monster)

Function which_monster():

def which_monster():
    while True:
        screen.blit(wybor_postaci, (0, 0))
        pygame.display.update()
        for event in pygame.event.get():
            if event.type == QUIT:
                pygame.quit()
                quit()
            if event.type == KEYDOWN:
                if event.key == K_1:
                    return monster1
                elif event.key == K_2:
                    return monster2
                elif event.key == K_3:
                    return monster3
                elif event.key == K_4:
                    return monster4
                elif event.key == K_5:
                    return monster5
                elif event.key == K_6:
                    return monster6

My question is: How I can get that pygame.Surface from which_monster function and assign it to variable?

I thought about something like that:

img = button("START!", 350, 350, 300, 50, blue, blue_light, which_monster)

or:

button("START!", 350, 350, 300, 50, blue, blue_light, img = which_monster)

but it's not working :(

Please help!