BSP Dungeon Generator

bspdungeon.py

import random
import pygame

class Node:
    def __init__(self, data=None):
        self.data = data
        self.left = None
        self.right = None


DUNGEON_WIDTH = 100
DUNGEON_HEIGHT = 100

MIN_ROOM_SIZE = 5
MIN_ROOM_RATIO = 0.45

def random_split(rect):
    if random.randint(0, 1) == 0:
        # Horizontal split
        r1 = rect.copy()
        r1.height = random.randrange(MIN_ROOM_SIZE, r1.height)
        r2 = rect.copy()
        r2.height -= r1.height
        r2.top += r1.height
        # Discard by ratio
        r1_h_ratio = r1.height / r1.width
        r2_h_ratio = r2.height / r2.width
        if r1_h_ratio < MIN_ROOM_RATIO or r2_h_ratio < MIN_ROOM_RATIO:
            return random_split(rect)
    else:
        # Vertical split
        # Horizontal split
        r1 = rect.copy()
        r1.width = random.randrange(MIN_ROOM_SIZE, r1.width)
        r2 = rect.copy()
        r2.width -= r1.width
        r2.left += r1.width
        # Discard by ratio
        r1_w_ratio = r1.width / r1.height
        r2_w_ratio = r2.width / r2.height
        if r1_w_ratio < MIN_ROOM_RATIO or r2_w_ratio < MIN_ROOM_RATIO:
            return random_split(rect)
    return r1, r2


def split_container(rect, niter):
    root = Node(data=rect)
    if niter != 0:
        left, right = random_split(rect)
        root.left = split_container(left, niter - 1)
        root.right = split_container(right, niter - 1)
    return root


def generate_dungeon():
    main_rect = pygame.Rect((0, 0), (DUNGEON_WIDTH, DUNGEON_HEIGHT))
    rect_tree = split_container(main_rect, 5)

    return rect_tree


def rect_list_from_tree(root_node):
    rect_list = []
    if root_node:
        rect_list = rect_list_from_tree(root_node.left)
        rect_list.append(root_node)
        rect_list.extend(rect_list_from_tree(root_node.right))
    return rect_list


def leaf_nodes_from_tree(root_node):
    node_list = []
    if root_node.left and root_node.right:
        node_list.extend(leaf_nodes_from_tree(root_node.left))
        node_list.extend(leaf_nodes_from_tree(root_node.right))
    else:
        node_list = [root_node]
    return node_list


def draw_paths(tree, surface):
    if tree.left is None or tree.right is None:
        return

    p1 = (tree.left.data.centerx * 5, tree.left.data.centery * 5)
    p2 = (tree.right.data.centerx * 5, tree.right.data.centery * 5)

    pygame.draw.line(surface, (127, 127, 127), p1, p2, width=10)

    draw_paths(tree.left, surface)
    draw_paths(tree.right, surface)


def draw_dungeon(dungeon_tree, surface):
    surface.fill((0, 0, 0))
    for node in rect_list_from_tree(dungeon_tree):
        r = node.data.copy()
        r.x *= 5
        r.y *= 5
        r.w *= 5
        r.h *= 5
        pygame.draw.rect(surface, (127, 127, 0), r, 1)

    rooms = []
    for node in leaf_nodes_from_tree(dungeon_tree):
        container_rect = node.data

        room_rect = container_rect.copy()
        room_rect.left = room_rect.left + random.randrange(0, room_rect.width // 2)
        room_rect.top = room_rect.top + random.randrange(0, room_rect.height // 2)
        room_rect.width = container_rect.width - (room_rect.left - container_rect.left)
        room_rect.height = container_rect.height - (room_rect.top - container_rect.top)
        room_rect.width -= random.randrange(0, room_rect.width // 2)
        room_rect.height -= random.randrange(0, room_rect.height // 2)
        rooms.append(room_rect)

        # Paint
        display_rect = room_rect.copy()
        display_rect.x *= 5
        display_rect.y *= 5
        display_rect.w *= 5
        display_rect.h *= 5
        pygame.draw.rect(surface, (127, 127, 127), display_rect)

    draw_paths(dungeon_tree, surface)
    pygame.display.flip()


def game_loop():
    pygame.init()
    screen = pygame.display.set_mode((DUNGEON_WIDTH * 5, DUNGEON_HEIGHT * 5))

    dungeon_tree = generate_dungeon()
    draw_dungeon(dungeon_tree, screen)

    clock = pygame.time.Clock()
    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                exit()
            if event.type == pygame.KEYUP and event.key == pygame.K_SPACE:
                dungeon_tree = generate_dungeon()
                draw_dungeon(dungeon_tree, screen)

        clock.tick(30)


if __name__ == "__main__":
    game_loop()