Coding problem for Pygame and higher-order functions