Toy program to make little nethack-style dungeons exported as tikz figures

generateBSPDungeon.py

#!/usr/bin/env python3
import random

BUFFER = 2

# For exporting the map as a tikz diagram
PREAMBLE = r'''\documentclass[tikz,convert=pdf2svg]{standalone}
\usepackage[utf8]{inputenc}
\usetikzlibrary{backgrounds}

\begin{document}
\tikzstyle{box} = [text centered, text width=0.5cm, draw=black, fill=yellow!30, anchor=south west]
\begin{tikzpicture}[]
'''

ENDTEX = r'''\end{tikzpicture}
\end{document}
'''

# A rectangle consists of a position, size, and list of other rectangles it has
# outgoing edges to
class Rectangle:
	def __init__(self, x, y, width, height):
		self.x = x
		self.y = y
		self.width = width
		self.height = height
		self.name = None
		self.edges = []
	def setName(self, name):
		self.name = name
	def addEdge(self, other):
		self.edges.append(other)
	def __str__(self):
		if( self.name == None ):
			return "[%dx%d rectangle at %d,%d]" % (self.width,self.height,self.x,self.y)
		else:
			return "[%dx%d rectangle at %d,%d named '%s']" % (self.width,self.height,self.x,self.y,self.name)
	def __repr__(self):
		return str(self)

# A BSP solution consists of a list of rectangles and their edges,
# and has utility functions to export to a tikz diagram
class BSP:
	def __init__(self, rectangles):
		self.rectangles = rectangles
		self.rectanglesByName = dict()
		for i,r in enumerate(self.rectangles):
			name = "r" + str(i+1)
			r.setName(name)
			self.rectanglesByName[name] = r
		self.edges = []
		for r in self.rectangles:
			for neighbor in r.edges:
				self.edges.append((r.name,neighbor.name))
	def __str__(self):
		rects = "Rectangles: " + str(self.rectangles) + "\n"
		edges = "Edges: " + str(self.edges)
		return rects + edges
	def __repr__(self):
		return str(self)
	def toTikz(self,filename):
		with open(filename, "w") as f:
			f.write(PREAMBLE)
			for r in self.rectangles:
				f.write("\\node (%s) at (%dmm,%dmm) [box,minimum width=%dmm,minimum height=%dmm] {%s};\n" % (r.name,r.x,r.y,r.width,r.height,r.name))
			f.write("\\begin{pgfonlayer}{background}")
			for e in self.edges:
				f.write("\\draw[-] (%s) -| (%s);\n" % e)
			f.write("\\end{pgfonlayer}")
			f.write(ENDTEX)

def generateBSP(width, height, minimum_width_ratio=0.1, minimum_height_ratio=0.1):
	minwidth = int(minimum_width_ratio * width)
	minheight = int(minimum_height_ratio * height)
	# Create a random room within a region
	def generateRectangle(x, y, width, height, mw, mh):
		w = random.randint(mw, width)
		h = random.randint(mh, height)
		maxx = (x + width) - w
		maxy = (y + height) - h
		rx = random.randint(x, maxx)
		ry = random.randint(y, maxy)
		return Rectangle(rx,ry,w,h)
		#return Rectangle(x-BUFFER,y-BUFFER,width+(2*BUFFER),height+(2*BUFFER))
	# Sub-divide a rectangle
	def generateCuts(x, y, width, height, mw, mh):
		print("Generating cuts [%d,%d %dx%d]" % (x,y,width,height))
		if( width <= 2*mw or height <= 2*mh ):
			return [generateRectangle(x+BUFFER, y+BUFFER, width-BUFFER, height-BUFFER, mw-(2*BUFFER), mh-(2*BUFFER))]
		# Choose vertical or horizontal slice
		cut1, cut2 = None, None
		direction = random.randint(0,1)
		if( direction == 0 ):
			minx = x + mw
			maxx = (x+width) - mw
			cutx = random.randint(minx, maxx)
			leftx,rightx = (x, cutx)
			leftw,rightw = (cutx - x, width - (cutx - x))
			cut1 = generateCuts(leftx, y, leftw, height, mw, mh)
			cut2 = generateCuts(rightx, y, rightw, height, mw, mh)
		else:
			miny = y + mh
			maxy = (y+height) - mh
			cuty = random.randint(miny, maxy)
			topy,bottomy = (y, cuty)
			toph,bottomh = (cuty - y, height - (cuty - y))
			cut1 = generateCuts(x, topy, width, toph, mw, mh)
			cut2 = generateCuts(x, bottomy, width, bottomh, mw, mh)
		# Pick a random rectangle from each side and add an edge between them
		r1 = random.choice(cut1)
		r2 = random.choice(cut2)
		r1.addEdge(r2)
		return cut1 + cut2
	rects = generateCuts(0, 0, width, height, minwidth, minheight)
	return BSP(rects)

if __name__ == "__main__":
	random.seed(17) # To produce the diagram from the blog post
	bsp = generateBSP(100,100)
	bsp.toTikz("foo.tex")