Pygame Snake, Pt. 3

Part 2 left off with a square moving on a grid. Now, let's take control of it by consuming KEYDOWN events.

One simple approach would be to have a variable that represents one of the four directions the snake can go. It might use values 0-3, or it could use string literals like "up" and "left".

A slightly more advanced technique would be to use the Vector2 class. We use pygame's Vector2 to keep track of the position of out snake part. We could also use a Vector2 to represent the snake's speed and direction. For example, the direction right would be represented by Vector2(1, 0), because if we add that vector to the current position, it will increase the x value (the first coordinate) by 1.

Left would be Vector2(-1, 0). Up and down would be Vector2(0, -1) and Vector2(0, 1), respectively. If up and down seem backwards, it's because in pygame (which uses a common computer drawing convention) the Y-axis increases as you go down, instead of up. The origin is the canvas's upper-left corner.

Add a new variable under dot:

vel = pygame.Vector2(1, 0)

Inside the for loop that consumes events, add the following under the first if statement:

        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_RIGHT:
                vel = pygame.Vector2(1, 0)
            if event.key == pygame.K_DOWN:
                vel = pygame.Vector2(0, 1)
            if event.key == pygame.K_LEFT:
                vel = pygame.Vector2(-1, 0)
            if event.key == pygame.K_UP:
                vel = pygame.Vector2(0, -1)

Finally, update the line of code that moves dot:

dot = dot + vel # or dot += vel

When you run your program now, you should be able to control the dot's movement with the arrow keys.

Full code at this point:

import pygame

W = 30
H = 30
S = 20

# pygame setup
pygame.init()
screen = pygame.display.set_mode((W * S, H * S))
clock = pygame.time.Clock()
running = True

dot = pygame.Vector2(W / 2, H / 2)
vel = pygame.Vector2(1, 0)

while running:
    # poll for events
    for event in pygame.event.get():
        # pygame.QUIT = user closed window
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_RIGHT:
                vel = pygame.Vector2(1, 0)
            if event.key == pygame.K_DOWN:
                vel = pygame.Vector2(0, 1)
            if event.key == pygame.K_LEFT:
                vel = pygame.Vector2(-1, 0)
            if event.key == pygame.K_UP:
                vel = pygame.Vector2(0, -1)

    # fill buffer with white
    screen.fill("white")

    dot += vel
    square = pygame.Rect(dot * S, (S, S))
    screen.fill("black", square)

    # copy buffer to screen
    pygame.display.flip()

    # limits FPS
    clock.tick(20)

pygame.quit()

OK, that's fine, but we could really tighten it up by introducing a dictionary to hold the key/velocity associations. Like this:

key_vel = {
    pygame.K_RIGHT: pygame.Vector2(1, 0),
    pygame.K_DOWN: pygame.Vector2(0, 1),
    pygame.K_LEFT: pygame.Vector2(-1, 0),
    pygame.K_UP: pygame.Vector2(0, -1)
}

Add that just above the while loop. Now the whole KEYDOWN if statement can be replaced with this:

        if event.type == pygame.KEYDOWN and event.key in key_vel:
            vel = key_vel[event.key]

When in is used with a dictionary like this, it checks the dictionary's keys to see if they contain the given value.

Full code now:

import pygame

W = 30
H = 30
S = 20

# pygame setup
pygame.init()
screen = pygame.display.set_mode((W * S, H * S))
clock = pygame.time.Clock()
running = True

dot = pygame.Vector2(W / 2, H / 2)
vel = pygame.Vector2(1, 0)

key_vel = {
    pygame.K_RIGHT: pygame.Vector2(1, 0),
    pygame.K_DOWN: pygame.Vector2(0, 1),
    pygame.K_LEFT: pygame.Vector2(-1, 0),
    pygame.K_UP: pygame.Vector2(0, -1)
}

while running:
    # poll for events
    for event in pygame.event.get():
        # pygame.QUIT = user closed window
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.KEYDOWN and event.key in key_vel:
            vel = key_vel[event.key]

    # fill buffer with white
    screen.fill("white")

    dot += vel
    square = pygame.Rect(dot * S, (S, S))
    screen.fill("black", square)

    # copy buffer to screen
    pygame.display.flip()

    # limits FPS
    clock.tick(20)

pygame.quit()