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

• my_stack_exchange_account: I had one idea for O(1) but placement where: you have one list containing empty spaces, one list containing occupied spaces, and one list containing where each coordinate in referenced in either list. Then when bugs are placed and the snake is moving, you switch coordinates around between both lists. But I'm not sure if that idea will actually work because I haven't tested it or thought about it very hard.

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