Snake implemented in Pygame

snake_main.py

import random
import pygame
from pygame.locals import *

class GameGrid:

    def __init__(self, grid_size=(100, 100), window_size=(400, 400)):
        self.grid_size = grid_size
        self.square_x_size = window_size[0] / grid_size[0]
        self.square_y_size = window_size[1] / grid_size[1]
        self.window = pygame.display.set_mode(window_size)
        self.grid_rects = self.get_grid_rects(grid_size)

        self.create_images()

        self.spawn_snake()
        self.spawn_apple()
        self.draw()
        self.pause()

    def create_images(self):
        self.background = pygame.Surface(self.window.get_size())
        self.background = self.background.convert()
        self.background.fill((0, 0, 0))

        self.snake_tile = pygame.Surface(
            (self.square_x_size, self.square_y_size))
        self.snake_tile.fill((255, 255, 255))

        self.apple_tile = pygame.Surface(
            (self.square_x_size, self.square_y_size))
        self.apple_tile.fill((255, 0, 0))

    def spawn_snake(self, pos=None):
        if pos is None:
            pos = (self.grid_size[0] / 2, self.grid_size[1] / 2)
        self.snake = Snake(self.grid_size)

    def draw(self):
        self.window.blit(self.background, (0, 0))
        for x, y in self.snake.body_squares:
            self.window.blit(self.snake_tile, self.grid_rects[x, y])
        self.window.blit(self.apple_tile, self.grid_rects[self.apple_pos])
        pygame.display.flip()

    def exit(self):
        pygame.display.quit()

    def get_grid_rects(self, grid_size):
        grid_rects = {}
        for row_num in range(grid_size[0]):
            for col_num in range(grid_size[1]):
                x_start = row_num * self.square_x_size
                x_end = (row_num + 1) * self.square_x_size
                y_start = col_num * self.square_y_size
                y_end = (col_num + 1) * self.square_y_size
                grid_rects[(row_num, col_num)] = (x_start, y_start, 
                                                  x_end, y_end)
        return grid_rects

    def main_loop(self):
        apple_spawn_time = 0
        while True:
            snake_ate_apple = False
            current_time = pygame.time.get_ticks()
            if self.snake.head_pos == self.apple_pos:
                snake_ate_apple = True
                self.spawn_apple()
                apple_spawn_time = current_time
            if current_time - apple_spawn_time > (5 * 1000):
                self.spawn_apple()
                apple_spawn_time = current_time
            self.process_events()
            snake_is_dead = self.snake.move(snake_ate_apple)
            self.draw()
            pygame.time.wait(40)
            if snake_is_dead:
                self.spawn_snake()
                self.pause()

    def pause(self):
        while True:
            for event in pygame.event.get():
                if event.type == QUIT:
                    self.exit()
                if event.type == KEYDOWN:
                    if event.key == K_ESCAPE:
                        self.exit()
                    if event.key == K_SPACE:
                        return


    def spawn_apple(self):
        random_x = random.choice(range(self.grid_size[0]))
        random_y = random.choice(range(self.grid_size[1]))
        self.apple_pos = (random_x, random_y)


    def process_events(self):
        for event in pygame.event.get():
            if event.type == QUIT:
                self.exit()
            if event.type == KEYDOWN:
                if event.key == K_ESCAPE:
                    self.exit()
                if event.key == K_SPACE:
                    self.pause()
                if event.key == K_UP:
                    if not self.snake.y_speed == 1:
                        self.snake.x_speed = 0
                        self.snake.y_speed = -1
                if event.key == K_DOWN:
                    if not self.snake.y_speed == -1:
                        self.snake.x_speed = 0
                        self.snake.y_speed = 1
                if event.key == K_LEFT:
                    if not self.snake.x_speed == 1:
                        self.snake.x_speed = -1
                        self.snake.y_speed = 0
                if event.key == K_RIGHT:
                    if not self.snake.x_speed == -1:
                        self.snake.x_speed = 1
                        self.snake.y_speed = 0



class Snake:

    def __init__(self, grid_size, length=10, head_pos=(12, 12)):
        self.grid_size = grid_size
        self.length = length
        self.head_pos = head_pos

        x, y = head_pos
        self.body_squares = [(x - i, y) for i in range(length)]

        self.x_speed = 1
        self.y_speed = 0

    def move(self, ate_apple=False):
        new_x = self.head_pos[0] + self.x_speed
        if new_x >= self.grid_size[0]:
            new_x = 1
        elif new_x == -1:
            new_x = self.grid_size[0] - 1
        new_y = self.head_pos[1] + self.y_speed
        if new_y >= self.grid_size[1]:
            new_y = 1
        elif new_y == -1:
            new_y = self.grid_size[1] - 1
        self.head_pos = (new_x, new_y)
        if self.head_pos in self.body_squares:
            dead = True
        else:
            dead = False
        if ate_apple:
            self.body_squares = [self.head_pos] + self.body_squares
        else:
            self.body_squares = [self.head_pos] + self.body_squares[:-1]
        return dead

def main():
    grid = GameGrid(grid_size=(30, 30))
    grid.main_loop()

if __name__ == '__main__':
    main()