Python Snake Game

import sys
import pygame
import random

# Screen dimensions
WIDTH, HEIGHT = 640, 480
CELL_SIZE = 20

# Colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GREEN = (0, 255, 0)
RED = (255, 0, 0)


class Snake:
    def __init__(self):
        self.segments = [(CELL_SIZE, CELL_SIZE), (CELL_SIZE * 2, CELL_SIZE)]
        self.direction = (CELL_SIZE, 0)

    def move(self):
        head = self.segments[0]
        new_head = (head[0] + self.direction[0], head[1] + self.direction[1])
        self.segments.insert(0, new_head)
        self.segments.pop()

    def grow(self):
        tail = self.segments[-1]
        new_tail = (tail[0] - self.direction[0], tail[1] - self.direction[1])
        self.segments.append(new_tail)

    def collide_with_itself(self):
        return self.segments[0] in self.segments[1:]

    def collide_with_bounds(self):
        head = self.segments[0]
        return head[0] < 0 or head[0] >= WIDTH or head[1] < 0 or head[1] >= HEIGHT


class Food:
    def __init__(self):
        self.position = self.random_position()

    def random_position(self):
        x = random.randrange(0, WIDTH, CELL_SIZE)
        y = random.randrange(0, HEIGHT, CELL_SIZE)
        return (x, y)


def main():
    pygame.init()
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    pygame.display.set_caption("Snake Game")

    snake = Snake()
    food = Food()

    clock = pygame.time.Clock()

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()

            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_UP:
                    snake.direction = (0, -CELL_SIZE)
                elif event.key == pygame.K_DOWN:
                    snake.direction = (0, CELL_SIZE)
                elif event.key == pygame.K_LEFT:
                    snake.direction = (-CELL_SIZE, 0)
                elif event.key == pygame.K_RIGHT:
                    snake.direction = (CELL_SIZE, 0)

        snake.move()

        if snake.collide_with_bounds() or snake.collide_with_itself():
            snake = Snake()
            food = Food()

        if snake.segments[0] == food.position:
            snake.grow()
            food.position = food.random_position()

        screen.fill(BLACK)

        for segment in snake.segments:
            pygame.draw.rect(screen, GREEN, pygame.Rect(segment[0], segment[1], CELL_SIZE, CELL_SIZE))

        pygame.draw.rect(screen, RED, pygame.Rect(food.position[0], food.position[1], CELL_SIZE, CELL_SIZE))

        pygame.display.flip()
        clock.tick(10)


if __name__ == "__main__":
    main()