26
\$\begingroup\$

I wrote a simple Python snake game which is about 250 lines of code. Can someone give me some advice on how I can refactor/make it better?

game.py

# game.py - 3/22/2013

import pygame, sys, os
from pygame.locals import *
from classes import *


def main():
    pygame.init()
    pygame.display.set_caption('PyGame Snake')

    window = pygame.display.set_mode((480, 480))
    screen = pygame.display.get_surface()
    clock = pygame.time.Clock()
    font = pygame.font.Font('freesansbold.ttf', 20)

    game = SnakeGame(window, screen, clock, font)

    while game.run(pygame.event.get()):
        pass

    pygame.quit()
    sys.exit()


if __name__ == '__main__':
    main()

classes.py

#classes.py - 3/22/2013

import pygame, random
from pygame.locals import *


# Gams speed
STARTING_FPS = 4
FPS_INCREMENT_FREQUENCY = 80


# Direction constants
DIRECTION_UP    = 1
DIRECTON_DOWN   = 2
DIRECTION_LEFT  = 3
DIRECTION_RIGHT = 4


# World size
WORLD_SIZE_X = 20
WORLD_SIZE_Y = 20


# Snake and food attributes
SNAKE_START_LENGTH = 4
SNAKE_COLOR = (0, 255, 0)
FOOD_COLOR = (255, 0, 0)


# Snake class
class Snake:

    # Initializes a Snake object
    def __init__(self, x, y, startLength):
        self.startLength = startLength
        self.startX = x
        self.startY = y
        self.reset()

    # Resets snake back to its original state
    def reset(self):
        self.pieces = []
        self.direction = 1

        for n in range(0, self.startLength):
            self.pieces.append((self.startX, self.startY + n))

    # Changes the direction of the snake
    def changeDirection(self, direction):
        # Moving in the opposite direction of current movement is not allowed
        if self.direction == 1 and direction == 2: return
        if self.direction == 2 and direction == 1: return
        if self.direction == 3 and direction == 4: return
        if self.direction == 4 and direction == 3: return

        self.direction = direction

    # Returns the head piece of the snake
    def getHead(self):
        return self.pieces[0]

    # Returns the tail piece of the snake
    def getTail(self):
        return self.pieces[len(self.pieces) - 1]

    # Updates snake by moving blocks in direction of movement
    def update(self):
        (headX, headY) = self.getHead()
        head = ()

        # Create new piece that is the new head of the snake
        if self.direction == 1: head = (headX, headY - 1)
        elif self.direction == 2: head = (headX, headY + 1)
        elif self.direction == 3: head = (headX - 1, headY)
        elif self.direction == 4: head = (headX + 1, headY)

        # Remove tail of the snake and add a new head
        self.pieces.insert(0, head)
        self.pieces.pop()

    # Adds a new piece to the end of the snake
    def grow(self):
        (tx, ty) = self.getTail()
        piece = ()

        if self.direction == 1: piece = (tx, ty + 1)
        elif self.direction == 2: piece = (tx, ty - 1)
        elif self.direction == 3: piece = (tx + 1, ty)
        elif self.direction == 4: piece = (tx - 1, ty)

        self.pieces.append(piece)

    # Are two pieces of the snake occupying the same block?
    def collidesWithSelf(self):
        """
        # Because of the way new pieces are added when the snake grows, eating a
        # new food block could cause the snake to die if it's in a certain position. 
        # So instead of checking if any of the spots have two pieces at once, the new
        # algorithm only checks if the position of the head piece contains more than one block.

        for p in self.pieces:
            if len(self.pieces) - len([c for c in self.pieces if c != p]) > 1: return True
        return False
        """

        return len([p for p in self.pieces if p == self.getHead()]) > 1


# SnakeGame class
class SnakeGame:

    # Initializes SnakeGame object with pre-initialized objects and configuration settings
    def __init__(self, window, screen, clock, font):
        self.window = window
        self.screen = screen
        self.clock = clock
        self.font = font

        self.fps = STARTING_FPS
        self.ticks = 0
        self.playing = True
        self.score = 0

        self.nextDirection = DIRECTION_UP
        self.sizeX = WORLD_SIZE_X
        self.sizeY = WORLD_SIZE_Y
        self.food = []
        self.snake = Snake(WORLD_SIZE_X / 2, WORLD_SIZE_Y / 2, SNAKE_START_LENGTH)

        self.addFood()

    # Adds a new piece of food to a random block
    def addFood(self):
        fx = None
        fy = None

        while fx is None or fy is None or (fx, fy) in self.food:
            fx = random.randint(1, self.sizeX)
            fy = random.randint(1, self.sizeY)

        self.food.append((fx, fy))

    # Handles input from keyboard
    def input(self, events):
        for e in events:
            if e.type == QUIT:
                return False

            elif e.type == KEYUP:
                if   e.key == K_w: self.nextDirection = 1
                elif e.key == K_s: self.nextDirection = 2
                elif e.key == K_a: self.nextDirection = 3
                elif e.key == K_d: self.nextDirection = 4
                elif e.key == K_SPACE and not self.playing: 
                    self.reset()

        return True

    # Update gamestate -- update snake and check for death
    def update(self):
        self.snake.changeDirection(self.nextDirection)
        self.snake.update()

        # If snake hits a food block, then consume the food, add new food and grow the snake
        for food in self.food: 
            if self.snake.getHead() == food:
                self.food.remove(food)
                self.addFood()
                self.snake.grow()
                self.score += len(self.snake.pieces) * 50

        # If snake collides with self or the screen boundaries, then game over
        (hx, hy) = self.snake.getHead()
        if self.snake.collidesWithSelf() or hx < 1 or hy < 1 or hx > self.sizeX or hy > self.sizeY:
            self.playing = False

    # Resets the game
    def reset(self):
        self.playing = True
        self.nextDirection = DIRECTION_UP
        self.fps = STARTING_FPS
        self.score = 0
        self.snake.reset()

    # Draws snake and food objects to the screen
    def draw(self):
        self.screen.fill((45, 45, 45))

        (width, height) = self.window.get_size()
        blockWidth = int(width / self.sizeX)
        blockHeight = int(height / self.sizeY)

        # Draw pieces of snake
        for (px, py) in self.snake.pieces: 
            pygame.draw.rect(self.screen, SNAKE_COLOR, (blockWidth * (px-1), blockHeight * (py-1), blockWidth, blockHeight))

        # Draw food objects
        for (fx, fy) in self.food:
            pygame.draw.rect(self.screen, FOOD_COLOR, (blockWidth * (fx-1), blockHeight * (fy-1), blockWidth, blockHeight))

        pygame.display.flip()

    # Draws the death message to the screen
    def drawDeath(self):
        self.screen.fill((255, 0, 0))
        self.screen.blit(self.font.render("Game over! Press Space to start a new game", 1, (255, 255, 255)), (20, 150))
        self.screen.blit(self.font.render("Your score is: %d" % self.score, 1, (255, 255, 255)), (140, 180))
        pygame.display.flip()

    # Run the main game loop
    def run(self, events):
        if not self.input(events): return False

        if self.playing: 
            self.update()
            self.draw()
        else: self.drawDeath()

        self.clock.tick(self.fps)

        self.ticks += 1
        if self.ticks % FPS_INCREMENT_FREQUENCY == 0: self.fps += 1

        return True
\$\endgroup\$

1 Answer 1

28
\$\begingroup\$

1. Introduction

This is not bad overall, considering that this is your first program written with PyGame. I've made many comments below, but don't take the length of this answer to heart: there are always many things to say about a piece of code of this length.

2. Game design issues

  1. The game could do with some instructions. I had to look at the source code to see that I need to use WASD for movement. Alternatively, you might allow the player to use the arrow keys too (these are natural keys the player might try).

  2. You use FPS (the number of frames per second) to control the speed of the snake. This design decision commits you to processing everything in the game at the same frequency as the snake moves. The concepts frames per second and speed of the snake in moves per second are distinct, so it's good practice to separate them.

    At the moment there's nothing in the game other than the snake, so you get away with this. But as soon as you add other game elements that need to animate at different speeds, you'll run up against this difficulty. Better to get this right while things are still simple.

    See section 5 for one way to solve this problem.

  3. New pieces of food can be created in positions occupied by the snake!

  4. The food position is not reset when a new game starts. (This can cause the snake to be overlap the food at the start of the game.)

  5. The score doesn't get drawn during the game.

3. Major comments

  1. The docstring for collidesWithSelf reads like this:

    """
    # Because of the way new pieces are added when the snake grows, eating a
    # new food block could cause the snake to die if it's in a certain position. 
    # So instead of checking if any of the spots have two pieces at once, the new
    # algorithm only checks if the position of the head piece contains more than one block.
    
    for p in self.pieces:
        if len(self.pieces) - len([c for c in self.pieces if c != p]) > 1: return True
    return False
    """
    

    This is not appropriate content for a docstring. The purpose of a docstring is to explain the interface of a method to a programmer who is trying to use it. But here you have some notes to yourself about the history of this function and why it is implemented like it is. These notes properly belong in a comment.

    The reason you have been having problems in this function is that the growth of the snake is not right. In the grow() method you grow a new tail segment in the opposite direction to the snake's current movement. But this can cause the snake to self-intersect.

    The usual way that "snake" games work is that when the snake eats some food, it does not grow a new tail segment immediately. Instead, it waits until the next time it moves and grows a new tail segment in the position where its old tail used to be. This is easily implemented by incrementing a counter each time the snake eats food:

    def grow(self):
        self.growth_pending += 1
    

    and then decrementing the counter instead of deleting the tail segment:

    if self.growth_pending > 0:
        self.growth_pending -= 1
    else:
        # Remove tail
        self.pieces.pop()
    

    This avoids self-intersection, and so this would allow you to implement the collision operation using your original approach. But you might consider this simpler approach:

    it = iter(self.pieces)
    head = next(it)
    return head in it
    
  2. You represent directions by numbers between 1 and 4. It is hard to remember which direction is which, so it would be easy to make a mistake and treat 1 as "up" in one part of the code but "down" in another. You'd be much less likely to make this mistake if you used the names DIRECTION_UP and so on. You went to all the trouble to create these names: why not use them?

    (But see 3.4 below for a better suggestion.)

  3. The code below looks dodgy because there is no else: on the end of the series of tests.

    head = ()
    if self.direction == 1: head = (headX, headY - 1)
    elif self.direction == 2: head = (headX, headY + 1)
    elif self.direction == 3: head = (headX - 1, headY)
    elif self.direction == 4: head = (headX + 1, headY)
    

    A programmer reading this would want to know what happens if self.direction is not in the range 1 to 4. Of course, you hope that you have designed the program so that this can't happen. So you might make this crystal clear by rewriting this code like this:

    if   self.direction == DIRECTION_UP:    head = (headX, headY - 1)
    elif self.direction == DIRECTION_DOWN:  head = (headX, headY + 1)
    elif self.direction == DIRECTION_LEFT:  head = (headX - 1, headY)
    elif self.direction == DIRECTION_RIGHT: head = (headX + 1, headY)
    else: raise RuntimeError("Bad direction: {}".format(self.direction))
    

    (But see 3.4 below for a better suggestion.)

  4. Instead of representing a direction with a number from 1 to 4 (which it's hard to remember which is which), why not use a pair (δx, δy)? For example, you could write:

    DIRECTION_UP    =  0, -1
    DIRECTION_DOWN  =  0,  1
    DIRECTION_LEFT  = -1,  0
    DIRECTION_RIGHT =  1,  0
    

    This would save you a bunch of if ... elif ... tests. For example you could rewrite the code I gave above like this:

    old_head = self.getHead()
    new_head = old_head[0] + self.direction[0], old_head[1] + self.direction[1]
    
  5. Similarly, instead of having a bunch of if .. elif ... tests for the input:

    if   e.key == K_w: self.nextDirection = 1
    elif e.key == K_s: self.nextDirection = 2
    elif e.key == K_a: self.nextDirection = 3
    elif e.key == K_d: self.nextDirection = 4
    

    you could have a lookup table that maps key to direction:

    # Input mapping
    KEY_DIRECTION = {
        K_w: DIRECTION_UP,    K_UP:    DIRECTION_UP,   
        K_s: DIRECTION_DOWN,  K_DOWN:  DIRECTION_DOWN, 
        K_a: DIRECTION_LEFT,  K_LEFT:  DIRECTION_LEFT, 
        K_d: DIRECTION_RIGHT, K_RIGHT: DIRECTION_RIGHT,
    }
    

    and then the test becomes:

    if e.key in KEY_DIRECTION:
        self.next_direction = KEY_DIRECTION[e.key]
    
  6. You check for collision with the edges of the playing area like this:

    (hx, hy) = self.snake.getHead()
    if hx < 1 or hy < 1 or hx > self.sizeX or hy > self.sizeY:
    

    but PyGame defines a Rect class with a collidepoint method, so you could set up the rectangle like this in SnakeGame.__init__:

    self.world = Rect(1, 1, 21, 21)
    

    and then test for collision like this:

    if not self.world.collidepoint(self.snake.getHead()):
    
  7. You number your coordinates starting at 1. This causes you some difficulty. For example, you have to subtract 1 from each of your coordinates before passing them to pygame.draw.rect. If your coordinates started at 0 you wouldn't have to do this.

  8. A lot of your code deals with positions or vectors represented by a pair (x, y). This means that each time you process a position, you have to break it down into its x and y coordinates, process the coordinates, and then reassemble the results into a new pair. For example:

    (headX, headY) = self.getHead()
    if self.direction == 1: head = (headX, headY - 1)
    

    You could simplify this code by making a class to represent positions and vectors. Sadly, PyGame does not come with such a class, but you can easily find many implementations on the web, or just write one yourself:

    class Vector(tuple):
        def __add__(self, other): return Vector(v + w for v, w in zip(self, other))
    

    now you can write:

    new_head = self.getHead() + self.direction
    

    and make many other simplifications. (See section 5 for more methods on the Vector class.)

4. Minor comments

  1. The Python style guide (PEP8) says that method names should "Use the function naming rules: lowercase with words separated by underscores as necessary to improve readability". So you should consider renaming collidesWithSelf as collides_with_self and so on. (You're not obliged to follow PEP8 but it makes it easier for other Python programmers to read your code.)

  2. The organization of the code into the files game.py and classes.py doesn't seem to be motivated by any principle, and the vague name classes.py confirms this. For a small game like this, I don't think there's anything to be lost by putting all the code in one file. (And if it grows to the point where you want to split it, the obvious thing to do would be to put the Snake class in its own module.)

  3. It's not clear to me what you gain from separating your game initialization code into main and SnakeGame.__init__. Why not put all the initialization into the latter?

  4. You don't need sys.exit() at the end of main: Python quits automatically when it finishes running your program. Adding this line just makes it hard to test your program from the interactive interpreter, because when you quit from the game, it exits the interactive interpreter too.

  5. You have unnecessary parentheses in many places. A line like:

    (headX, headY) = self.getHead()
    

    can be written:

    headX, headY = self.getHead()
    

    since comma binds more tighly than assignment in Python. "Program as if you know the language"!

  6. DIRECTON_DOWN is misspelled. (Is this why you don't use it?)

  7. You represent the snake using a queue of positions, which is a good approach. However, you implement your queue using a Python list. The trouble here is that Python lists are efficient at adding and removing elements at the end, but not at the beginning. In particular the operation

    self.pieces.insert(0, head)
    

    takes time proportional to the length of self.pieces. (You can consult the TimeComplexity page on the Python wiki to see the time complexity of operations on built-in Python data structures.) This isn't a big deal here, since the snake never gets very long, but it's worth getting into practice at thinking about the complexity of your algorithms.

    For an efficient queue implementation, use collections.deque.

  8. The line:

    return self.pieces[len(self.pieces) - 1]
    

    can be rewritten:

    return self.pieces[-1]
    

    Since negative list indexes count backwards from the end of the list.

  9. Instead of doing division and coercing the result to an integer:

    int(width / self.sizeX)
    

    use Python's floor division operation:

    width // self.sizeX
    

5. Revised code

This code addresses the comments above and includes some more improvements for you to discover.

from collections import deque
import pygame
from random import randrange
import sys
from pygame.locals import *

class Vector(tuple):
    """A tuple that supports some vector operations.

    >>> v, w = Vector((1, 2)), Vector((3, 4))
    >>> v + w, w - v, v * 10, 100 * v, -v
    ((4, 6), (2, 2), (10, 20), (100, 200), (-1, -2))
    """
    def __add__(self, other): return Vector(v + w for v, w in zip(self, other))
    def __radd__(self, other): return Vector(w + v for v, w in zip(self, other))
    def __sub__(self, other): return Vector(v - w for v, w in zip(self, other))
    def __rsub__(self, other): return Vector(w - v for v, w in zip(self, other))
    def __mul__(self, s): return Vector(v * s for v in self)
    def __rmul__(self, s): return Vector(v * s for v in self)
    def __neg__(self): return -1 * self

FPS = 60                        # Game frames per second
SEGMENT_SCORE = 50              # Score per segment

SNAKE_SPEED_INITIAL = 4.0       # Initial snake speed (squares per second)
SNAKE_SPEED_INCREMENT = 0.25    # Snake speeds up this much each time it grows
SNAKE_START_LENGTH = 4          # Initial snake length in segments

WORLD_SIZE = Vector((20, 20))   # World size, in blocks
BLOCK_SIZE = 24                 # Block size, in pixels

BACKGROUND_COLOR = 45, 45, 45
SNAKE_COLOR = 0, 255, 0
FOOD_COLOR = 255, 0, 0
DEATH_COLOR = 255, 0, 0
TEXT_COLOR = 255, 255, 255

DIRECTION_UP    = Vector(( 0, -1))
DIRECTION_DOWN  = Vector(( 0,  1))
DIRECTION_LEFT  = Vector((-1,  0))
DIRECTION_RIGHT = Vector(( 1,  0))
DIRECTION_DR    = DIRECTION_DOWN + DIRECTION_RIGHT

# Map from PyGame key event to the corresponding direction.
KEY_DIRECTION = {
    K_w: DIRECTION_UP,    K_UP:    DIRECTION_UP,   
    K_s: DIRECTION_DOWN,  K_DOWN:  DIRECTION_DOWN, 
    K_a: DIRECTION_LEFT,  K_LEFT:  DIRECTION_LEFT, 
    K_d: DIRECTION_RIGHT, K_RIGHT: DIRECTION_RIGHT,
}

class Snake(object):
    def __init__(self, start, start_length):
        self.speed = SNAKE_SPEED_INITIAL # Speed in squares per second.
        self.timer = 1.0 / self.speed    # Time remaining to next movement.
        self.growth_pending = 0          # Number of segments still to grow.
        self.direction = DIRECTION_UP    # Current movement direction.
        self.segments = deque([start - self.direction * i
                             for i in xrange(start_length)])

    def __iter__(self):
        return iter(self.segments)

    def __len__(self):
        return len(self.segments)

    def change_direction(self, direction):
        """Update the direction of the snake."""
        # Moving in the opposite direction of current movement is not allowed.
        if self.direction != -direction:
            self.direction = direction

    def head(self):
        """Return the position of the snake's head."""
        return self.segments[0]

    def update(self, dt, direction):
        """Update the snake by dt seconds and possibly set direction."""
        self.timer -= dt
        if self.timer > 0:
            # Nothing to do yet.
            return
        # Moving in the opposite direction of current movement is not allowed.
        if self.direction != -direction:
            self.direction = direction
        self.timer += 1 / self.speed
        # Add a new head.
        self.segments.appendleft(self.head() + self.direction)
        if self.growth_pending > 0:
            self.growth_pending -= 1
        else:
            # Remove tail.
            self.segments.pop()

    def grow(self):
        """Grow snake by one segment and speed up."""
        self.growth_pending += 1
        self.speed += SNAKE_SPEED_INCREMENT

    def self_intersecting(self):
        """Is the snake currently self-intersecting?"""
        it = iter(self)
        head = next(it)
        return head in it

class SnakeGame(object):
    def __init__(self):
        pygame.display.set_caption('PyGame Snake')
        self.block_size = BLOCK_SIZE
        self.window = pygame.display.set_mode(WORLD_SIZE * self.block_size)
        self.screen = pygame.display.get_surface()
        self.clock = pygame.time.Clock()
        self.font = pygame.font.Font('freesansbold.ttf', 20)
        self.world = Rect((0, 0), WORLD_SIZE)
        self.reset()

    def reset(self):
        """Start a new game."""
        self.playing = True
        self.next_direction = DIRECTION_UP
        self.score = 0
        self.snake = Snake(self.world.center, SNAKE_START_LENGTH)
        self.food = set()
        self.add_food()

    def add_food(self):
        """Ensure that there is at least one piece of food.
        (And, with small probability, more than one.)
        """
        while not (self.food and randrange(4)):
            food = Vector(map(randrange, self.world.bottomright))
            if food not in self.food and food not in self.snake:
                self.food.add(food)

    def input(self, e):
        """Process keyboard event e."""
        if e.key in KEY_DIRECTION:
            self.next_direction = KEY_DIRECTION[e.key]
        elif e.key == K_SPACE and not self.playing:
            self.reset()

    def update(self, dt):
        """Update the game by dt seconds."""
        self.snake.update(dt, self.next_direction)

        # If snake hits a food block, then consume the food, add new
        # food and grow the snake.
        head = self.snake.head()
        if head in self.food:
            self.food.remove(head)
            self.add_food()
            self.snake.grow()
            self.score += len(self.snake) * SEGMENT_SCORE

        # If snake collides with self or the screen boundaries, then
        # it's game over.
        if self.snake.self_intersecting() or not self.world.collidepoint(self.snake.head()):
            self.playing = False

    def block(self, p):
        """Return the screen rectangle corresponding to the position p."""
        return Rect(p * self.block_size, DIRECTION_DR * self.block_size)

    def draw_text(self, text, p):
        """Draw text at position p."""
        self.screen.blit(self.font.render(text, 1, TEXT_COLOR), p)

    def draw(self):
        """Draw game (while playing)."""
        self.screen.fill(BACKGROUND_COLOR)
        for p in self.snake:
            pygame.draw.rect(self.screen, SNAKE_COLOR, self.block(p))
        for f in self.food:
            pygame.draw.rect(self.screen, FOOD_COLOR, self.block(f))
        self.draw_text("Score: {}".format(self.score), (20, 20))

    def draw_death(self):
        """Draw game (after game over)."""
        self.screen.fill(DEATH_COLOR)
        self.draw_text("Game over! Press Space to start a new game", (20, 150))
        self.draw_text("Your score is: {}".format(self.score), (140, 180))

    def play(self):
        """Play game until the QUIT event is received."""
        while True:
            dt = self.clock.tick(FPS) / 1000.0 # convert to seconds
            for e in pygame.event.get():
                if e.type == QUIT:
                    return
                elif e.type == KEYUP:
                    self.input(e)
            if self.playing: 
                self.update(dt)
                self.draw()
            else:
                self.draw_death()
            pygame.display.flip()

def main():
    pygame.init()
    SnakeGame().play()
    pygame.quit()
\$\endgroup\$
4
  • \$\begingroup\$ hi thank you, this code is very useful. However I wonder how to implement no world boundaries so when the snake exit the world it comes back at the opposite side. if not self.world.collidepoint(self.snake.head()): #reposition snake here. Any help would be very welcome because I don't know anything about vectors :-( \$\endgroup\$ Sep 11, 2014 at 12:40
  • \$\begingroup\$ @kasperTaeymans: That would make a good question for Stack Overflow, if you made your best attempt first, and then asked them for help with fixing your code. \$\endgroup\$ Sep 11, 2014 at 14:37
  • \$\begingroup\$ Hi Gareth, thank you for the reply! I will ask it on stackoverflow when I have time. I'm buzzy on a different part of my project now. I'll post a link here in the comments when the question + attempt is online. It might be useful for other interested readers. Also I'll link back to this code review. \$\endgroup\$ Sep 12, 2014 at 13:48
  • \$\begingroup\$ I've asked a question related to the no boundaries option on stackoverflow. Here is the link to that question: stackoverflow.com/questions/25891680/… \$\endgroup\$ Sep 17, 2014 at 13:17

Not the answer you're looking for? Browse other questions tagged or ask your own question.