Pygame Snake, Pt. 3

Source: Dev.to

Controlling the Square with Keyboard Input

Using a Direction Vector

A convenient way to store the snake’s speed and direction is with pygame.Vector2.
The four directions can be represented as:

Note: In Pygame the Y‑axis increases downward, so “up” uses a negative Y component.

Add a velocity variable after creating the dot:

vel = pygame.Vector2(1, 0)   # start moving right

Inside the event loop, handle the arrow keys:

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 dot’s position each frame:

dot = dot + vel   # or simply: dot += vel

Running the program now lets you move the square with the arrow keys.

Full Example (Step 1)

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():
        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)

    # clear screen
    screen.fill("white")

    # move and draw the square
    dot += vel
    square = pygame.Rect(dot * S, (S, S))
    screen.fill("black", square)

    # update display
    pygame.display.flip()
    clock.tick(20)

pygame.quit()

Refactoring with a Dictionary

The repeated if statements can be replaced by a dictionary that maps keys to velocity vectors:

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)
}

Place this definition just before the main loop. Then the KEYDOWN handling becomes:

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

Full Example (Refactored)

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)

# map keys to velocity vectors
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():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.KEYDOWN and event.key in key_vel:
            vel = key_vel[event.key]

    # clear screen
    screen.fill("white")

    # move and draw the square
    dot += vel
    square = pygame.Rect(dot * S, (S, S))
    screen.fill("black", square)

    # update display
    pygame.display.flip()
    clock.tick(20)

pygame.quit()

With the dictionary approach the code is shorter, clearer, and easier to extend (e.g., adding more controls).