mirror of
https://github.com/CHOMPStation2/CHOMPStation2.git
synced 2025-12-10 18:22:39 +00:00
460 lines
13 KiB
Python
460 lines
13 KiB
Python
# Tools for working with DreamMaker maps
|
|
|
|
import io
|
|
import bidict
|
|
import random
|
|
from collections import namedtuple
|
|
|
|
TGM_HEADER = "//MAP CONVERTED BY dmm2tgm.py THIS HEADER COMMENT PREVENTS RECONVERSION, DO NOT REMOVE"
|
|
ENCODING = 'utf-8'
|
|
|
|
Coordinate = namedtuple('Coordinate', ['x', 'y', 'z'])
|
|
|
|
class DMM:
|
|
__slots__ = ['key_length', 'size', 'dictionary', 'grid', 'header']
|
|
|
|
def __init__(self, key_length, size):
|
|
self.key_length = key_length
|
|
self.size = size
|
|
self.dictionary = bidict.bidict()
|
|
self.grid = {}
|
|
self.header = None
|
|
|
|
@staticmethod
|
|
def from_file(fname):
|
|
# stream the file rather than forcing all its contents to memory
|
|
with open(fname, 'r', encoding=ENCODING) as f:
|
|
return _parse(iter(lambda: f.read(1), ''))
|
|
|
|
@staticmethod
|
|
def from_bytes(bytes):
|
|
return _parse(bytes.decode(ENCODING))
|
|
|
|
def to_file(self, fname, tgm = False):
|
|
with open(fname, 'w', newline='\n', encoding=ENCODING) as f:
|
|
(save_tgm if tgm else save_dmm)(self, f)
|
|
|
|
def to_bytes(self, tgm = False):
|
|
bio = io.BytesIO()
|
|
with io.TextIOWrapper(bio, newline='\n', encoding=ENCODING) as f:
|
|
(save_tgm if tgm else save_dmm)(self, f)
|
|
f.flush()
|
|
return bio.getvalue()
|
|
|
|
def generate_new_key(self):
|
|
# ensure that free keys exist by increasing the key length if necessary
|
|
free_keys = (BASE ** self.key_length) - len(self.dictionary)
|
|
while free_keys <= 0:
|
|
self.key_length += 1
|
|
free_keys = (BASE ** self.key_length) - len(self.dictionary)
|
|
|
|
# choose one of the free keys at random
|
|
key = 0
|
|
while free_keys:
|
|
if key not in self.dictionary:
|
|
# this construction is used to avoid needing to construct the
|
|
# full set in order to random.choice() from it
|
|
if random.random() < 1 / free_keys:
|
|
return key
|
|
free_keys -= 1
|
|
key += 1
|
|
|
|
raise RuntimeError("ran out of keys, this shouldn't happen")
|
|
|
|
@property
|
|
def coords_zyx(self):
|
|
for z in range(1, self.size.z + 1):
|
|
for y in range(1, self.size.y + 1):
|
|
for x in range(1, self.size.x + 1):
|
|
yield (z, y, x)
|
|
|
|
@property
|
|
def coords_z(self):
|
|
return range(1, self.size.z + 1)
|
|
|
|
@property
|
|
def coords_yx(self):
|
|
for y in range(1, self.size.y + 1):
|
|
for x in range(1, self.size.x + 1):
|
|
yield (y, x)
|
|
|
|
# ----------
|
|
# key handling
|
|
|
|
# Base 52 a-z A-Z dictionary for fast conversion
|
|
BASE = 52
|
|
base52 = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
|
|
base52_r = {x: i for i, x in enumerate(base52)}
|
|
assert len(base52) == BASE and len(base52_r) == BASE
|
|
|
|
def key_to_num(key):
|
|
num = 0
|
|
for ch in key:
|
|
num = BASE * num + base52_r[ch]
|
|
return num
|
|
|
|
def num_to_key(num, key_length):
|
|
if num >= BASE ** key_length:
|
|
raise KeyTooLarge(f"num={num} does not fit in key_length={key_length}")
|
|
|
|
result = ''
|
|
while num:
|
|
result = base52[num % BASE] + result
|
|
num //= BASE
|
|
|
|
assert len(result) <= key_length
|
|
return base52[0] * (key_length - len(result)) + result
|
|
|
|
class KeyTooLarge(Exception):
|
|
pass
|
|
|
|
# ----------
|
|
# An actual atom parser
|
|
|
|
def parse_map_atom(atom):
|
|
try:
|
|
i = atom.index('{')
|
|
except ValueError:
|
|
return atom, {}
|
|
|
|
path, rest = atom[:i], atom[i+1:]
|
|
vars = {}
|
|
|
|
in_string = False
|
|
in_name = False
|
|
escaping = False
|
|
current_name = ''
|
|
current = ''
|
|
for ch in rest:
|
|
if escaping:
|
|
escaping = False
|
|
current += ch
|
|
elif ch == '\\':
|
|
escaping = True
|
|
elif ch == '"':
|
|
in_string = not in_string
|
|
current += ch
|
|
elif in_string:
|
|
current += ch
|
|
elif ch == ';':
|
|
vars[current_name.strip()] = current.strip()
|
|
current_name = current = ''
|
|
elif ch == '=':
|
|
current_name = current
|
|
current = ''
|
|
elif ch == '}':
|
|
vars[current_name.strip()] = current.strip()
|
|
break
|
|
elif ch not in ' ':
|
|
current += ch
|
|
|
|
return path, vars
|
|
|
|
# ----------
|
|
# TGM writer
|
|
|
|
def save_tgm(dmm, output):
|
|
output.write(f"{TGM_HEADER}\n")
|
|
if dmm.header:
|
|
output.write(f"{dmm.header}\n")
|
|
|
|
# write dictionary in tgm format
|
|
for key, value in sorted(dmm.dictionary.items()):
|
|
output.write(f'"{num_to_key(key, dmm.key_length)}" = (\n')
|
|
for idx, thing in enumerate(value):
|
|
in_quote_block = False
|
|
in_varedit_block = False
|
|
for char in thing:
|
|
if in_quote_block:
|
|
if char == '"':
|
|
in_quote_block = False
|
|
output.write(char)
|
|
elif char == '"':
|
|
in_quote_block = True
|
|
output.write(char)
|
|
elif not in_varedit_block:
|
|
if char == "{":
|
|
in_varedit_block = True
|
|
output.write("{\n\t")
|
|
else:
|
|
output.write(char)
|
|
elif char == ";":
|
|
output.write(";\n\t")
|
|
elif char == "}":
|
|
output.write("\n\t}")
|
|
in_varedit_block = False
|
|
else:
|
|
output.write(char)
|
|
if idx < len(value) - 1:
|
|
output.write(",\n")
|
|
output.write(")\n")
|
|
|
|
# thanks to YotaXP for finding out about this one
|
|
max_x, max_y, max_z = dmm.size
|
|
for z in range(1, max_z + 1):
|
|
output.write("\n")
|
|
for x in range(1, max_x + 1):
|
|
output.write(f"({x},{1},{z}) = {{\"\n")
|
|
for y in range(1, max_y + 1):
|
|
output.write(f"{num_to_key(dmm.grid[x, y, z], dmm.key_length)}\n")
|
|
output.write("\"}\n")
|
|
|
|
# ----------
|
|
# DMM writer
|
|
|
|
def save_dmm(dmm, output):
|
|
if dmm.header:
|
|
output.write(f"{dmm.header}\n")
|
|
|
|
# writes a tile dictionary the same way Dreammaker does
|
|
for key, value in sorted(dmm.dictionary.items()):
|
|
output.write(f'"{num_to_key(key, dmm.key_length)}" = ({",".join(value)})\n')
|
|
|
|
output.write("\n")
|
|
|
|
# writes a map grid the same way Dreammaker does
|
|
max_x, max_y, max_z = dmm.size
|
|
for z in range(1, max_z + 1):
|
|
output.write(f"(1,1,{z}) = {{\"\n")
|
|
|
|
for y in range(1, max_y + 1):
|
|
for x in range(1, max_x + 1):
|
|
try:
|
|
output.write(num_to_key(dmm.grid[x, y, z], dmm.key_length))
|
|
except KeyError:
|
|
print(f"Key error: ({x}, {y}, {z})")
|
|
output.write("\n")
|
|
output.write("\"}\n")
|
|
|
|
# ----------
|
|
# Parser
|
|
|
|
def _parse(map_raw_text):
|
|
in_comment_line = False
|
|
comment_trigger = False
|
|
|
|
in_quote_block = False
|
|
in_key_block = False
|
|
in_data_block = False
|
|
in_varedit_block = False
|
|
after_data_block = False
|
|
escaping = False
|
|
skip_whitespace = False
|
|
|
|
dictionary = bidict.bidict()
|
|
duplicate_keys = {}
|
|
curr_key_len = 0
|
|
curr_key = 0
|
|
curr_datum = ""
|
|
curr_data = list()
|
|
|
|
in_map_block = False
|
|
in_coord_block = False
|
|
in_map_string = False
|
|
iter_x = 0
|
|
adjust_y = True
|
|
|
|
curr_num = ""
|
|
reading_coord = "x"
|
|
|
|
key_length = 0
|
|
|
|
maxx = 0
|
|
maxy = 0
|
|
maxz = 0
|
|
|
|
curr_x = 0
|
|
curr_y = 0
|
|
curr_z = 0
|
|
grid = dict()
|
|
|
|
it = iter(map_raw_text)
|
|
|
|
# map block
|
|
for char in it:
|
|
if char == "\n":
|
|
in_comment_line = False
|
|
comment_trigger = False
|
|
continue
|
|
elif in_comment_line:
|
|
continue
|
|
elif char == "\t":
|
|
continue
|
|
|
|
if char == "/" and not in_quote_block:
|
|
if comment_trigger:
|
|
in_comment_line = True
|
|
continue
|
|
else:
|
|
comment_trigger = True
|
|
else:
|
|
comment_trigger = False
|
|
|
|
if in_data_block:
|
|
|
|
if in_varedit_block:
|
|
|
|
if in_quote_block:
|
|
if char == "\\":
|
|
curr_datum = curr_datum + char
|
|
escaping = True
|
|
|
|
elif escaping:
|
|
curr_datum = curr_datum + char
|
|
escaping = False
|
|
|
|
elif char == "\"":
|
|
curr_datum = curr_datum + char
|
|
in_quote_block = False
|
|
|
|
else:
|
|
curr_datum = curr_datum + char
|
|
|
|
else:
|
|
if skip_whitespace and char == " ":
|
|
skip_whitespace = False
|
|
continue
|
|
skip_whitespace = False
|
|
|
|
if char == "\"":
|
|
curr_datum = curr_datum + char
|
|
in_quote_block = True
|
|
|
|
elif char == ";":
|
|
skip_whitespace = True
|
|
curr_datum = curr_datum + char
|
|
|
|
elif char == "}":
|
|
curr_datum = curr_datum + char
|
|
in_varedit_block = False
|
|
|
|
else:
|
|
curr_datum = curr_datum + char
|
|
|
|
elif char == "{":
|
|
curr_datum = curr_datum + char
|
|
in_varedit_block = True
|
|
|
|
elif char == ",":
|
|
curr_data.append(curr_datum)
|
|
curr_datum = ""
|
|
|
|
elif char == ")":
|
|
curr_data.append(curr_datum)
|
|
curr_data = tuple(curr_data)
|
|
try:
|
|
dictionary[curr_key] = curr_data
|
|
except bidict.ValueDuplicationError:
|
|
# if the map has duplicate values, eliminate them now
|
|
duplicate_keys[curr_key] = dictionary.inv[curr_data]
|
|
curr_data = list()
|
|
curr_datum = ""
|
|
curr_key = 0
|
|
curr_key_len = 0
|
|
in_data_block = False
|
|
after_data_block = True
|
|
|
|
else:
|
|
curr_datum = curr_datum + char
|
|
|
|
elif in_key_block:
|
|
if char == "\"":
|
|
in_key_block = False
|
|
if key_length == 0:
|
|
key_length = curr_key_len
|
|
else:
|
|
assert key_length == curr_key_len
|
|
else:
|
|
curr_key = BASE * curr_key + base52_r[char]
|
|
curr_key_len += 1
|
|
|
|
# else we're looking for a key block, a data block or the map block
|
|
elif char == "\"":
|
|
in_key_block = True
|
|
after_data_block = False
|
|
|
|
elif char == "(":
|
|
if after_data_block:
|
|
in_coord_block = True
|
|
after_data_block = False
|
|
curr_key = 0
|
|
curr_key_len = 0
|
|
break
|
|
else:
|
|
in_data_block = True
|
|
after_data_block = False
|
|
|
|
# grid block
|
|
for char in it:
|
|
if in_coord_block:
|
|
if char == ",":
|
|
if reading_coord == "x":
|
|
curr_x = int(curr_num)
|
|
if curr_x > maxx:
|
|
maxx = curr_x
|
|
iter_x = 0
|
|
curr_num = ""
|
|
reading_coord = "y"
|
|
elif reading_coord == "y":
|
|
curr_y = int(curr_num)
|
|
if curr_y > maxy:
|
|
maxy = curr_y
|
|
curr_num = ""
|
|
reading_coord = "z"
|
|
else:
|
|
raise ValueError("too many dimensions")
|
|
|
|
elif char == ")":
|
|
curr_z = int(curr_num)
|
|
if curr_z > maxz:
|
|
maxz = curr_z
|
|
in_coord_block = False
|
|
reading_coord = "x"
|
|
curr_num = ""
|
|
|
|
else:
|
|
curr_num = curr_num + char
|
|
|
|
elif in_map_string:
|
|
if char == "\"":
|
|
in_map_string = False
|
|
adjust_y = True
|
|
curr_y -= 1
|
|
|
|
elif char == "\n":
|
|
if adjust_y:
|
|
adjust_y = False
|
|
else:
|
|
curr_y += 1
|
|
if curr_x > maxx:
|
|
maxx = curr_x
|
|
if iter_x > 1:
|
|
curr_x = 1
|
|
iter_x = 0
|
|
|
|
else:
|
|
curr_key = BASE * curr_key + base52_r[char]
|
|
curr_key_len += 1
|
|
if curr_key_len == key_length:
|
|
iter_x += 1
|
|
if iter_x > 1:
|
|
curr_x += 1
|
|
|
|
grid[curr_x, curr_y, curr_z] = duplicate_keys.get(curr_key, curr_key)
|
|
curr_key = 0
|
|
curr_key_len = 0
|
|
|
|
# else look for coordinate block or a map string
|
|
elif char == "(":
|
|
in_coord_block = True
|
|
elif char == "\"":
|
|
in_map_string = True
|
|
|
|
if curr_y > maxy:
|
|
maxy = curr_y
|
|
|
|
data = DMM(key_length, Coordinate(maxx, maxy, maxz))
|
|
data.dictionary = dictionary
|
|
data.grid = grid
|
|
return data
|