python pygame snake game

Question

I used numpy for the O(n) parts, but I didn't add any ability to change grid size because I still wonder if there's a way to make it better than O(n).

main.py:

import pygame
import numpy as np
from random import randrange
from clickablebox import ClickableBox

# values of grid spaces
EMPTY = 0
BUG = 1
SNAKE_UP = 2
SNAKE_RIGHT = 3
SNAKE_DOWN = 4
SNAKE_LEFT = 5

GRID_COLORS = (
    (200, 200, 200),
    (65, 163, 23),
    (255, 243, 128),
    (255, 243, 128),
    (255, 243, 128),
    (255, 243, 128)
)

GRID_Y_OFFSET = 20
GRID_X_OFFSET = 20
GRID_HEIGHT = 16
GRID_WIDTH = 16
GRID_SPACE_SIZE = 35

TICKS_BETWEEN_MOVEMENT = 15 # 1 tick is 16 milliseconds

class Game:
    def __init__(self):
        # numpy used instead of list for self.available_bug_spaces because it needs to find and to remove values, which is O(n)
        self.available_bug_spaces = None
        
        self.grid = []
        self.tail = (0, 0)
        self.head = (0, 0)
    
    
    def setup_game(self):
        # setting up lists
        self.available_bug_spaces = np.arange(GRID_HEIGHT * GRID_WIDTH, dtype=np.uint64)
        self.grid.clear()
        for _ in range(GRID_HEIGHT):
            self.grid.append([EMPTY] * GRID_WIDTH)
        
        # setting initial snake spaces
        mid_y = GRID_HEIGHT // 2
        mid_x = GRID_WIDTH // 2
        for x in range(mid_x, mid_x + 5):
            self.grid[mid_y][x] = SNAKE_LEFT
        
        self.available_bug_spaces = np.delete(self.available_bug_spaces, range((mid_y * GRID_WIDTH) + mid_x, (mid_y * GRID_WIDTH) + mid_x + 5))
        self.place_bug()
        
        self.head = (mid_y, mid_x)
        self.tail = (mid_y, x)
    
    
    def place_bug(self):
        # I can't figure out how to make this part better than O(n)
        if len(self.available_bug_spaces) == 0:
            return (-1, -1) # player won
        available_bug_spaces_index = randrange(len(self.available_bug_spaces))
        available_bug_spaces_value = self.available_bug_spaces[available_bug_spaces_index]
        bug_y, bug_x = int(available_bug_spaces_value) // GRID_WIDTH, int(available_bug_spaces_value) % GRID_WIDTH
        self.available_bug_spaces = np.delete(self.available_bug_spaces, (available_bug_spaces_index,))
        self.grid[bug_y][bug_x] = BUG
        return (bug_y, bug_x)
    
    
    def get_next_coord(self, coord, direction):
        vertical_change, horizontal_change = ((-1, 0), (0, 1), (1, 0), (0, -1))[direction - SNAKE_UP]
        return (coord[0] + vertical_change, coord[1] + horizontal_change)
    
    
    def move_forward_no_bug(self, next_head_coord, head_direction):
        old_tail_coord = self.tail
        tail_direction = self.grid[self.tail[0]][self.tail[1]]
        next_tail_coord = self.get_next_coord(self.tail, tail_direction)
        self.grid[self.tail[0]][self.tail[1]] = 0
        self.grid[self.head[0]][self.head[1]] = head_direction
        self.grid[next_head_coord[0]][next_head_coord[1]] = head_direction
        
        # switch out new head position with old tail position
        # I can't figure out how to make this part better than O(n)
        
        self.available_bug_spaces[np.where(self.available_bug_spaces == (next_head_coord[0] * GRID_WIDTH) + next_head_coord[1])[0][0]] = (self.tail[0] * GRID_WIDTH) + self.tail[1]
        
        self.head = next_head_coord
        self.tail = next_tail_coord
        
        return self.head, old_tail_coord
    
    
    def move_forward_eat_bug(self, next_head_coord, head_direction):
        self.grid[self.head[0]][self.head[1]] = head_direction
        self.grid[next_head_coord[0]][next_head_coord[1]] = head_direction
        bug_coord = self.place_bug()
        
        self.head = next_head_coord
        
        return self.head, bug_coord


class GUI:
    
    def __init__(self):
        self.win = pygame.display.set_mode((1200, 700))
        self.font = pygame.font.Font(pygame.font.get_default_font(), 30)
        self.start_button = ClickableBox(self.win, self.font, "start", (1105, 515), (1090, 480), (100, 100))
        self.quit_button = ClickableBox(self.win, self.font, "quit", (1110, 625), (1090, 590), (100, 100))
        self.start_button.draw()
        self.quit_button.draw()
        self.menu_drawn = True
    
    
    def toggle_menu_screen(self):
        self.set_message("")
        if self.menu_drawn: # this means the menu is currently drawn
            self.start_button.undraw()
            self.menu_drawn = False
        else:
            pygame.draw.rect(self.win, (0, 0, 0), (GRID_X_OFFSET, GRID_Y_OFFSET, GRID_SPACE_SIZE * GRID_WIDTH, GRID_SPACE_SIZE * GRID_HEIGHT)) # undrawing grid
            self.start_button.draw()
            self.menu_drawn = True
    
    
    def set_message(self, message):
        pygame.draw.rect(self.win, (0, 0, 0), (150, 600, 940, 100))
        text_surface = self.font.render(message, True, (50, 50, 50))
        self.win.blit(text_surface, (150, 650))
    
    
    def draw_grid(self, grid):
        for y in range(GRID_HEIGHT):
            for x in range(GRID_WIDTH):
                pygame.draw.rect(self.win, GRID_COLORS[grid[y][x]], (GRID_X_OFFSET + (x * GRID_SPACE_SIZE), GRID_Y_OFFSET + (y * GRID_SPACE_SIZE), GRID_SPACE_SIZE, GRID_SPACE_SIZE))
    
    
    def draw_coords(self, grid, coords):
        for coord in coords:
            pygame.draw.rect(self.win, GRID_COLORS[grid[coord[0]][coord[1]]], (GRID_X_OFFSET + (coord[1] * GRID_SPACE_SIZE), GRID_Y_OFFSET + (coord[0] * GRID_SPACE_SIZE), GRID_SPACE_SIZE, GRID_SPACE_SIZE))


def wait_for_unpause(gui):
    while True:
        pygame.time.wait(16)
        pygame.display.update()
        
        events = pygame.event.get()
        for event in events:
            if event.type == pygame.QUIT:
                return False, False
            elif event.type == pygame.MOUSEBUTTONUP and event.button == 1 and gui.quit_button.clicked(pygame.mouse.get_pos()):
                return False, True
            elif event.type == pygame.KEYDOWN and event.key == pygame.K_p:
                return True, True


def wait_for_input(gui):
    while True:
        pygame.time.wait(16)
        pygame.display.update()
        
        events = pygame.event.get()
        for event in events:
            if (
                (event.type == pygame.MOUSEBUTTONUP and event.button == 1 and gui.quit_button.clicked(pygame.mouse.get_pos()))
                or event.type == pygame.KEYDOWN
                or event.type == pygame.QUIT
                ):
                return event.type != pygame.MOUSEBUTTONUP, pygame.MOUSEBUTTONUP != pygame.QUIT


def play_game(game, gui):
    gui.set_message("")
    ticks_since_last_movement = 0
    direction = None
    current_head_direction = game.grid[game.head[0]][game.head[1]]
    
    while True:
        pygame.time.wait(16)
        pygame.display.update()
        ticks_since_last_movement += 1
        
        events = pygame.event.get()
        for event in events:
            if event.type == pygame.QUIT:
                return False, False
            elif event.type == pygame.MOUSEBUTTONUP and event.button == 1 and gui.quit_button.clicked(pygame.mouse.get_pos()):
                return False, True
            elif event.type == pygame.KEYDOWN:
                if (event.key == pygame.K_UP or event.key == pygame.K_w) and current_head_direction != SNAKE_DOWN:
                    direction = SNAKE_UP
                elif (event.key == pygame.K_RIGHT or event.key == pygame.K_d) and current_head_direction != SNAKE_LEFT:
                    direction = SNAKE_RIGHT
                elif (event.key == pygame.K_DOWN or event.key == pygame.K_s) and current_head_direction != SNAKE_UP:
                    direction = SNAKE_DOWN
                elif (event.key == pygame.K_LEFT or event.key == pygame.K_a) and current_head_direction != SNAKE_RIGHT:
                    direction = SNAKE_LEFT
                elif event.key == pygame.K_p:
                    ticks_since_last_movement = 0
                    gui.set_message("paused")
                    keep_playing, keep_running = wait_for_unpause(gui)
                    if not keep_playing or not keep_running:
                        return keep_playing, keep_running
                    gui.set_message("")
        if direction is not None or ticks_since_last_movement == TICKS_BETWEEN_MOVEMENT:
            if direction is None:
                direction = game.grid[game.head[0]][game.head[1]]
            ticks_since_last_movement = 0
            next_head_y, next_head_x = game.get_next_coord(game.head, direction)
            if not 0 <= next_head_y < GRID_HEIGHT or not 0 <= next_head_x < GRID_WIDTH or SNAKE_UP <= game.grid[next_head_y][next_head_x] <= SNAKE_LEFT:
                gui.set_message("you lose, press any key to reset")
                return wait_for_input(gui)
            if game.grid[next_head_y][next_head_x] == EMPTY:
                head, old_tail = game.move_forward_no_bug((next_head_y, next_head_x), direction)
                gui.draw_coords(game.grid, (head, old_tail))
            else:
                head, bug_coord = game.move_forward_eat_bug((next_head_y, next_head_x), direction)
                if bug_coord == (-1, -1):
                    gui.set_message("you win, press any key to reset")
                    return wait_for_input(gui)
                gui.draw_coords(game.grid, (head, bug_coord))
            current_head_direction = game.grid[game.head[0]][game.head[1]]
            direction = None


def main():
    pygame.init()
    pygame.display.set_caption("snake game")
    game = Game()
    gui = GUI()
    
    keep_running = True
    while keep_running:
        pygame.time.wait(16)
        pygame.display.update()
        
        events = pygame.event.get()
        for event in events:
            if event.type == pygame.QUIT:
                keep_running = False
            elif event.type == pygame.MOUSEBUTTONUP and event.button == 1:
                if gui.quit_button.clicked(pygame.mouse.get_pos()):
                    keep_running = False
                elif gui.start_button.clicked(pygame.mouse.get_pos()):
                    gui.toggle_menu_screen()
                    keep_playing = True
                    while keep_playing:
                        game.setup_game()
                        gui.set_message("press any key to start")
                        gui.draw_grid(game.grid)
                        pygame.display.update()
                        keep_playing, keep_running = wait_for_input(gui)
                        if keep_playing:
                            keep_playing, keep_running = play_game(game, gui)
                    gui.toggle_menu_screen()
    
    pygame.quit()


if __name__ == "__main__":
    main()

clickablebox.py (TextEntryBox not used in this):

import pygame

class ClickableBox:
    
    def __init__(self, win, font, text, text_location, box_location, dimensions):
        self.font = font
        self.text = text
        self.text_location = text_location
        self.box_location = box_location
        self.dimensions = dimensions
        
        self.win = win
        self.box = None
    
    
    def draw(self, color=(150, 150, 150)):
        self.box = pygame.draw.rect(self.win, color, (*self.box_location, *self.dimensions))
        text_surface = self.font.render(self.text, True, (50, 50, 50))
        self.win.blit(text_surface, self.text_location)
    
    
    def undraw(self):
        pygame.draw.rect(self.win, (0, 0, 0), (*self.box_location, *self.dimensions))
        self.box = None
    
    
    def clicked(self, coords):
        return self.box is not None and self.box.collidepoint(coords)


class TextEntryBox(ClickableBox):
    
    def __init__(self, win, font, text, text_location, box_location, dimensions, typing_location, max_text):
        super().__init__(win, font, text, text_location, box_location, dimensions)
        self.typing_location = typing_location
        self.max_text = max_text
        self.box_text = ""
    
    
    def draw(self, text, color=(150, 150, 150)):
        super().draw(color)
        text = str(text)[:self.max_text]
        typing_surface = self.font.render(text, True, (70, 70, 70))
        self.win.blit(typing_surface, self.typing_location)
        self.box_text = text
    
    
    def undraw(self):
        upper_left_undraw_box = (self.box_location[0], self.text_location[1])
        undraw_box_dimensions = (self.dimensions[0], self.dimensions[1] + (self.box_location[1] - self.text_location[1]))
        pygame.draw.rect(self.win, (0, 0, 0), (*upper_left_undraw_box, *undraw_box_dimensions))
        self.box = None

Answer 1

Notice that, directly after construction, Game is unusable and requires a call to setup_game. You also have one Game instance that's reused via multiple setup_game calls that each trample the contents of the class. To fix these issues, and to deter harmful state mutation, on the inside of your main() loop just construct a game instance each time. Merge your setup code into the constructor.

Add PEP484 type hints to your function signatures.

I can't figure out how to make this part better than O(n)

The O(n) comes from your np.delete. I'm not convinced that there's a very good role for Numpy in this application, especially for available_bug_spaces.

The first instinct is to use a built-in set to represent your available_bug_spaces; this will have amortised O(1) membership test, insertion and deletion. But this comes with an asterisk: if you try to random.sample() on the set, this is now deprecated; if you try it now you will see something like

DeprecationWarning: Sampling from a set deprecated
since Python 3.9 and will be removed in a subsequent version.

Read e.g. random.choice from set? and its links.

Since your grid is only 16x16, efficiency is not that big a deal, so it's almost certainly not worth implementing an exotic data structure for this purpose. You could do worse than to

• Use a built-in set of available coordinates
• When moving normally, all operations will be O(1) amortised
• When placing a new bug, cast to a tuple to random.choice() on it; this will be O(n) but is relatively rare so it makes sense to eat the cost here