Step 6 — Functions

def draw_game(screen, snake, food, score, font)

All drawing code in one place. Takes the screen surface, game state, and font. Returns nothing (implicitly returns None). The function side-effects the screen directly.

def draw_game(screen, snake, food, score, font):
    """Draw the current game state to the screen."""
    screen.fill(BG)
    # Draw snake body
    for i, (cx, cy) in enumerate(snake):
        colour = DKGREEN if i == 0 else GREEN
        pygame.draw.rect(screen, colour,
            pygame.Rect(cx*CELL_SIZE, cy*CELL_SIZE,
                        CELL_SIZE-2, CELL_SIZE-2), border_radius=4)
    # Draw food
    fx, fy = food
    pygame.draw.circle(screen, RED,
        (fx*CELL_SIZE + CELL_SIZE//2,
         fy*CELL_SIZE + CELL_SIZE//2), CELL_SIZE//3)
    # Draw score
    screen.blit(font.render(f"Score: {score}", True, WHITE), (8, 6))
    pygame.display.flip()

def update_game(snake, direction, food, score)

All game logic in one place. Returns the updated state — Python can return multiple values as a tuple.

def update_game(snake, direction, food, score):
    """Move snake, check food & collisions. Returns (snake, food, score, running)."""
    hx, hy = snake[0]
    if direction == "RIGHT": hx += 1
    elif direction == "LEFT": hx -= 1
    elif direction == "UP":   hy -= 1
    elif direction == "DOWN": hy += 1

    new_head = (hx, hy)
    if hx < 0 or hx >= GRID_COLS or hy < 0 or hy >= GRID_ROWS:
        return snake, food, score, False   # game over

    ate_food = (new_head == food)
    snake.insert(0, new_head)
    if ate_food:
        score += 1
        food = (random.randint(0,GRID_COLS-1),
                random.randint(0,GRID_ROWS-1))
    else:
        snake.pop()

    if new_head in snake[1:]:
        return snake, food, score, False   # game over

    return snake, food, score, True   # all good

# Usage:
snake, food, score, running = update_game(snake, direction, food, score)

def game_over_screen(screen, score, font)

Shows a "Game Over" screen and waits for a key. Extracted so it can be called cleanly after the while loop.

def game_over_screen(screen, score, font):
    """Show game over message and wait for key press."""
    screen.fill(BG)
    msg1 = font.render("GAME OVER", True, RED)
    msg2 = font.render(f"Score: {score} — Press any key", True, WHITE)
    cx, cy = screen.get_width()//2, screen.get_height()//2
    screen.blit(msg1, msg1.get_rect(center=(cx, cy-24)))
    screen.blit(msg2, msg2.get_rect(center=(cx, cy+16)))
    pygame.display.flip()
    waiting = True
    while waiting:
        for event in pygame.event.get():
            if event.type in (pygame.QUIT, pygame.KEYDOWN):
                waiting = False

Add a draw_grid function 💡 Abstraction

Create def draw_grid(screen): that draws faint lines between cells. Use pygame.draw.line(screen, (50,50,50), (x,0), (x,WINDOW_H)) in a for loop for columns, then rows.

Restart on game over 📋 Algorithm

In game_over_screen, return True if the user pressed a key (restart) and False for QUIT. Wrap the entire game in a while True: and check the return: if not game_over_screen(...): break.

Move direction handling to a function 🔄 Pattern

Create def handle_key(key, current_direction): that returns the new direction string. Call it from the event loop: direction = handle_key(event.key, direction).