Simple pygame based snake game.

snake.py

from __future__ import annotations
import random
from collections import deque
import time
import colorsys

import pygame

WIDTH, HEIGHT = 50, 50
TILE_WIDTH = 10
TILE_HEIGHT = 10


def blit_grid(surf: pygame.Surface, snake: deque, pellet, direction):
    food = pygame.Surface((TILE_WIDTH, TILE_HEIGHT))
    food.fill('green')

    surf.fill('black')
    surf.blit(food, (pellet[0] * TILE_WIDTH, pellet[1] * TILE_HEIGHT))
    for i, (x, y) in enumerate(snake):
        rgb = colorsys.hls_to_rgb(1, .5 + (i / len(snake) / 2), 1)
        r, g, b = (int(x * 255) for x in rgb)
        tile = pygame.Surface((TILE_WIDTH, TILE_HEIGHT))
        tile.fill((r, g, b))

        if i == 0:
            if direction == 'up':
                pygame.draw.rect(tile, (255, 255, 255), (1, 1, 2, 2))
                pygame.draw.rect(tile, (255, 255, 255), (6, 1, 2, 2))
            elif direction == 'down':
                pygame.draw.rect(tile, (255, 255, 255), (1, 6, 2, 2))
                pygame.draw.rect(tile, (255, 255, 255), (6, 6, 2, 2))
            elif direction == 'left':
                pygame.draw.rect(tile, (255, 255, 255), (1, 1, 2, 2))
                pygame.draw.rect(tile, (255, 255, 255), (1, 6, 2, 2))
            elif direction == 'right':
                pygame.draw.rect(tile, (255, 255, 255), (6, 1, 2, 2))
                pygame.draw.rect(tile, (255, 255, 255), (6, 6, 2, 2))

        surf.blit(tile, (x * TILE_WIDTH, y * TILE_HEIGHT))


def new_pellet(snake: deque):
    coords = set(snake)
    while True:
        pellet = random.randrange(WIDTH), random.randrange(HEIGHT)
        if pellet not in coords:
            return pellet


def main():
    pygame.init()
    screen = pygame.display.set_mode((WIDTH * TILE_WIDTH, HEIGHT * TILE_HEIGHT))
    clock = pygame.time.Clock()

    start_length = 1
    snake = deque(maxlen=start_length)
    for _ in range(start_length):
        snake.append((WIDTH // 2, HEIGHT // 2))
    direction = 'up'
    pellet = new_pellet(snake)

    time_passed = 0
    tick_ms = 100
    paused = False

    while True:
        time_passed += clock.tick(60)
        if time_passed < tick_ms:
            continue
        time_passed -= tick_ms

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                return
            if event.type == pygame.KEYDOWN:
                if not paused:
                    if event.key == pygame.K_UP:
                        direction = 'up'
                    elif event.key == pygame.K_DOWN:
                        direction = 'down'
                    elif event.key == pygame.K_LEFT:
                        direction = 'left'
                    elif event.key == pygame.K_RIGHT:
                        direction = 'right'

                if event.key == pygame.K_ESCAPE:
                    pygame.quit()
                    return
                elif event.key == pygame.K_SPACE:
                    paused = not paused

        if paused:
            continue

        hx, hy = snake[0]

        if direction == 'up':
            nx, ny = hx, hy - 1
        elif direction == 'down':
            nx, ny = hx, hy + 1
        elif direction == 'left':
            nx, ny = hx - 1, hy
        elif direction == 'right':
            nx, ny = hx + 1, hy
        else:
            raise ValueError(f'Invalid value for direction: {direction}')

        if (nx, ny) == pellet:
            snake = deque(snake, maxlen=snake.maxlen + 1)
            pellet = new_pellet(snake)
        elif (nx, ny) in snake:
            break
        elif nx < 0 or nx >= HEIGHT or ny < 0 or ny >= WIDTH:
            break

        snake.appendleft((nx, ny))

        blit_grid(screen, snake, pellet, direction)
        pygame.display.flip()

    font = pygame.font.SysFont('Arial', 30)
    game_over = font.render('GAME OVER', True, 'red')

    x = (screen.get_width() - game_over.get_width()) // 2
    y = (screen.get_height() - game_over.get_height()) // 2
    start = time.time()
    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                return
            if event.type == pygame.KEYUP:
                # just a small 1 sec delay before accepting key press events
                if time.time() - start < 1:
                    break
                pygame.quit()
                return

        screen.fill('black')
        screen.blit(game_over, (x, y))
        pygame.display.flip()
        clock.tick(30)


if __name__ == '__main__':
    main()