pathfinding example

pathfinder.py

"""
Functions for pygame pathfinding.
"""
import pygame as pg


class Cell(object):
    def __init__(self, index, cell_size, occupant=None): 
        self.index = index
        self.occupant = occupant
        topleft = index[0] * cell_size, index[1] * cell_size
        self.rect = pg.Rect(topleft, (cell_size, cell_size))
        
    def get_neighbors(self, grid):
        neighbors = set()
        offsets = ((-1, 0), (1, 0), (0, -1), (0, 1))
        for offset in offsets:
            x, y = offset[0] + self.index[0], offset[1] + self.index[1]
            try:
                if grid[(x, y)].occupant is None:            
                    neighbors.add((x,y))
            except KeyError:
                pass
        return neighbors


def make_grid(num_columns, num_rows, cell_size):
    """Make a dict of index: Cell pairs."""
    grid = {} 
    for col in range(num_columns):
        for row in range(num_rows):
            grid[(col, row)] = Cell((col, row), cell_size)
    return grid    
    
    
def find_path(origin, destination, grid, depth):
    paths = get_paths(origin, destination, grid, depth)
    return back_trace(paths, origin, destination)

    
def back_trace(candidates, origin, destination):
    """Work back from the destination to the origin."""
    backwards = candidates[::-1]
    path = []
    target = destination
    for candidate in backwards:
        for c, came_from in candidate:
            if c == origin:
                path.append(came_from)
                return path[::-1]
            if c == target:
                path.append(c)
                target = came_from
                break
    return path[::-1]

    
def get_paths(origin, destination, grid, depth):    
    #localize variables to cut down on lookups
    origin = origin
    dest = destination
    grid = grid
    
    #cells visited so far
    visited = {origin}
    
    #A list of sets, each set will consist of 2-tuples of grid-index tuples.
    #The first element of the tuple is the grid index travelled to, the second
    #is the index travelled from. Each set added represents the next 
    #"depth level" in the search. The last set in the list will end up
    #containing the destination.    
    candidates = [{(origin, origin)}]
    
    for d in range(depth):
        new_candidates = set()
        for c, came_from in candidates[d]:
            new_neighbors = grid[c].get_neighbors(grid)
            for new_neighbor in new_neighbors:
                if new_neighbor not in visited:
                    visited.add(new_neighbor)
                    new_candidates.add((new_neighbor, c))
                if new_neighbor == dest:
                    candidates.append(new_candidates)
                    return candidates                 
        candidates.append(new_candidates)
    return candidates

pathfinding.py

"""
Sinple pygame program to test out the
pathfinder module.
"""

import sys
from random import randint
import pygame as pg
import pathfinder as pf


class Rock(object):
    def __init__(self, index, rect):
        self.index = index
        self.rect = rect
        
    def draw(self, surface):
        pg.draw.rect(surface, pg.Color("gray40"), self.rect)

        
class Ship(object):
    def __init__(self, index, rect):
        self.index = index
        self.rect = rect
        
    def draw(self, surface):
        pg.draw.rect(surface, pg.Color("saddlebrown"), self.rect)

        
class Pathfinding(object):
    def __init__(self):
        self.done = False
        self.clock = pg.time.Clock()
        self.fps = 60
        num_columns = 60
        num_rows = 60
        cell_size = 10
        width = num_columns * cell_size
        height = num_rows * cell_size
        self.screen = pg.display.set_mode((width, height))
        self.grid = pf.make_grid(num_columns, num_rows, cell_size)
        self.ship = Ship((2, 2), self.grid[(2, 2)].rect)
        self.rocks = []
        for _ in range(40):
            spot = randint(0, num_columns - 1), randint(0, num_rows - 1)
            if spot != self.ship.index:
                self.place_rock(spot)
        self.path = []
    
    def place_rock(self, index):
        rock = Rock(index, self.grid[index].rect)
        self.grid[index].occupant = rock
        self.rocks.append(rock)
        
    def draw(self):
        self.screen.fill(pg.Color("darkblue"))
        for cell in self.grid:
            pg.draw.rect(self.screen, pg.Color("black"), self.grid[cell].rect, 2)
        for rock in self.rocks:
            rock.draw(self.screen)
        for i in self.path:
            pg.draw.rect(self.screen, pg.Color("lightblue"), self.grid[i].rect)
        self.ship.draw(self.screen)
        
    def event_loop(self):
        for event in pg.event.get():
            if event.type == pg.QUIT:
                self.done = True
            elif event.type == pg.MOUSEBUTTONDOWN:
                for idx in self.grid:
                    if self.grid[idx].rect.collidepoint(event.pos):
                        if event.button == 1:
                            self.path = pf.find_path(self.ship.index, idx, self.grid, 100)
                        elif event.button == 3:
                            self.ship.index = idx
                            self.ship.rect = self.grid[idx].rect
                            self.path = []
                            
    def update(self, dt):
        pass

    def game_loop(self):
        while not self.done:
            dt = self.clock.tick(self.fps)
            self.event_loop()
            self.update(dt)
            self.draw()
            pg.display.update()
        
        
if __name__ == "__main__":
    path_test = Pathfinding()
    path_test.game_loop()
    pg.quit()
    sys.exit()