maze.py

from PIL import Image, ImageDraw
import random


class Cell:
    wall_pairs = {'N': 'S', 'S': 'N', 'E': 'W', 'W': 'E'}

    def __init__(self, x, y):
        self.x, self.y = x, y
        self.walls = {'N': True, 'S': True, 'E': True, 'W': True}

    def has_all_walls(self):
        return all(self.walls.values())

    def knock_down_wall(self, other, wall):
        self.walls[wall] = False
        other.walls[Cell.wall_pairs[wall]] = False


class Maze:
    def __init__(self, nx, ny, ix=0, iy=0):
        self.nx, self.ny = nx, ny
        self.ix, self.iy = ix, iy
        self.maze_map = [[Cell(x, y) for y in range(ny)] for x in range(nx)]

    def cell_at(self, x, y):
        return self.maze_map[x][y]

    def find_valid_neighbours(self, cell):
        delta = [('W', (-1, 0)), ('E', (1, 0)), ('S', (0, 1)), ('N', (0, -1))]
        neighbours = []
        for direction, (dx, dy) in delta:
            x2, y2 = cell.x + dx, cell.y + dy
            if (0 <= x2 < self.nx) and (0 <= y2 < self.ny):
                neighbour = self.cell_at(x2, y2)
                if neighbour.has_all_walls():
                    neighbours.append((direction, neighbour))
        return neighbours

    def make_maze(self):
        n = self.nx * self.ny
        cell_stack = []
        current_cell = self.cell_at(self.ix, self.iy)
        nv = 1

        while nv < n:
            neighbours = self.find_valid_neighbours(current_cell)

            if not neighbours:
                current_cell = cell_stack.pop()
                continue

            direction, next_cell = random.choice(neighbours)
            current_cell.knock_down_wall(next_cell, direction)
            cell_stack.append(current_cell)
            current_cell = next_cell
            nv += 1

    def write_png(self, filename):
        cell_size = 20
        img_width = self.nx * cell_size
        img_height = self.ny * cell_size
        img = Image.new("RGB", (img_width, img_height), "white")
        draw = ImageDraw.Draw(img)

        for x in range(self.nx):
            for y in range(self.ny):
                cell = self.cell_at(x, y)
                x1, y1 = x * cell_size, y * cell_size
                x2, y2 = (x + 1) * cell_size, (y + 1) * cell_size

                if cell.walls['N']:
                    draw.line([(x1, y1), (x2, y1)], fill="black", width=1)
                if cell.walls['S']:
                    draw.line([(x1, y2), (x2, y2)], fill="black", width=1)
                if cell.walls['E']:
                    draw.line([(x2, y1), (x2, y2)], fill="black", width=1)
                if cell.walls['W']:
                    draw.line([(x1, y1), (x1, y2)], fill="black", width=1)

        img.save(filename)


if __name__ == '__main__':
    maze = Maze(10, 10)
    maze.make_maze()
    maze.write_png("maze_output.png")