Structure code

mt2obj.py

@@ -61,6 +61,8 @@ def convert(desc, vals):
 			out += int(val, 16),
 		elif c == 'b': # bool
 			out += str2bool(val),
+		elif c == 'A': # special: arglist
+			out += parse_arglist(val),
 		i += 1
 	return out
 
@@ -94,6 +96,32 @@ def parse_arglist(s):
 		out[tmp1] = tmp2
 	return out
 
+class MtsReader():
+	def __init__(self):
+		self.dim = (0, 0, 0) # W x H x D
+		self.namemap = {}
+		self.data = b""
+	def decode(self, f):
+		assert(f.mode == "rb")
+		if f.read(4) != b"MTSM":
+			raise Exception("Incorrect magic value, this isn't a schematic!")
+		ver = struct.unpack("!H", f.read(2))[0]
+		if ver not in (3, 4):
+			raise Exception("Wrong file version: got %d, expected 3 or 4" % v)#
+		self.dim = struct.unpack("!HHH", f.read(6))
+		f.seek(self.dim[1], 1) # skip some stuff
+		count = struct.unpack("!H", f.read(2))[0]
+		for i in range(count):
+			l = struct.unpack("!H", f.read(2))[0]
+			self.namemap[i] = f.read(l).decode("ascii")
+		self.data = zlib.decompress(f.read())
+	def dimensions(self):
+		return self.dim
+	def getnode(self, x, y, z):
+		off = (x + y*self.dim[0] + z*self.dim[0]*self.dim[1]) * 2
+		nid = struct.unpack("!H", self.data[off:off+2])[0]
+		return self.namemap[nid]
+
 class PreprocessingException(Exception):
 	pass
 
@@ -130,6 +158,8 @@ class Preprocessor():
 		for m in re.finditer(r'{([a-zA-Z0-9_-]+)}', line):
 			name = m.group(1)
 			if name not in self.vars.keys():
+				# TODO: this is currently required as undefined vars can
+				#       occur in ifdef sections (which are processed later)
 				tmp += line[last:m.start()] + "???"
 				last = m.end()
 				continue
@@ -150,7 +180,6 @@ class Preprocessor():
 			if inst == "define":
 				args = self._splitinto(args, " ", 2)
 				self.vars[args[0]] = args[1]
-				return None
 			elif inst == "if":
 				self._assertne(args)
 				self.state = 1 if str2bool(args) else 2
@@ -181,115 +210,93 @@ def usage():
 	print("  -n    Set path to nodes.txt")
 	exit(1)
 
-nodetbl = {}
-
-r_entry = re.compile(r'^(\S+) (\S+) (\d+) (\d+) (\d+) (\S+)(?: (\d+))?$')
 
 optargs, args = getopt.getopt(sys.argv[1:], 'tn:')
-
+if len(args) != 1:
-# TODO: structure code below
-if len(args) < 1:
 	usage()
 
 nodetbl = {}
+r_entry = re.compile(r'(\S+) (\S+) (\d+) (\d+) (\d+) (\S+)(?: (\d+))?')
 with open(getarg(optargs, "-n", default="nodes.txt"), "r") as f:
 	for line in f:
 		m = r_entry.match(line)
-		nodetbl[m.group(1)] = convert('siiisi', m.groups()[1:])
+		nodetbl[m.group(1)] = convert('siiiAi', m.groups()[1:])
 
-if True:
+schem = MtsReader()
-	sch = open(args[0], "rb")
+with open(args[0], "rb") as f:
-	if sch.read(4) != b"MTSM":
+	try:
-		print("This file does not look like an MTS schematic..")
+		schem.decode(f)
+	except Exception as e:
+		print(str(e), file=sys.stderr)
 		exit(1)
-	v = struct.unpack("!H", sch.read(2))[0]
+
-	if v not in (3, 4):
+filepart = ".".join(args[0].split(".")[:-1])
-		print("Wrong file version: got %d, expected 3 or 4" % v)
+objfile = filepart + ".obj"
-		exit(1)
+mtlfile = filepart + ".mtl"
-	width, height, depth = struct.unpack("!HHH", sch.read(6))
+
-	sch.seek(height, 1)
+nodes_seen = set()
-	nodecount = struct.unpack("!H", sch.read(2))[0]
+with open(objfile, "w") as obj:
-	nodemap = {}
+	obj.write("# Exported by mt2obj (https://github.com/sfan5/mt2obj)\nmtllib %s\n\n\n" % mtlfile)
-	for i in range(nodecount):
-		l = struct.unpack("!H", sch.read(2))[0]
-		name = sch.read(l).decode('ascii')
-		nodemap[i] = name
-	# TODO use zlib.decompressobj() instead of decompressing everything at once
-	cdata = sch.read()
-	sch.close()
-	data = zlib.decompress(cdata)
-	del cdata
-	filepart = args[0][:args[0].find(".")]
-	obj = open(filepart + ".obj", "w")
-	obj.write("# Exported by mt2obj\nmtllib %s\n\n\n" % (filepart + ".mtl", ))
 	i = 0
-	foundnodes = []
+	for x in range(schem.dimensions()[0]):
-	unknownnodes = []
+		for y in range(schem.dimensions()[1]):
-	for x in range(width):
+			for z in range(schem.dimensions()[2]):
-		for y in range(height):
+				node = schem.getnode(x, y, z)
-			for z in range(depth):
+				if node == "air":
-				off = (x + y*width + z*width*height) * 2
-				nid = struct.unpack("!H", data[off:off + 2])[0]
-				nname = nodemap[nid]
-				if nname == "air":
 					continue
-				if not nname in nodetbl.keys():
+				nodes_seen.add(node)
-					if not nname in unknownnodes:
+				if node not in nodetbl.keys():
-						unknownnodes.append(nname)
 					continue
-				else:
+				obj.write("o node%d\nusemtl %s\n" % (i, node.replace(":", "__")))
-					if not nname in foundnodes:
+				with open("models/%s.obj" % nodetbl[node][0], "r") as objdef:
-						foundnodes.append(nname)
+					pp = Preprocessor(omit_empty=True)
-				obj.write("o node%d\n" % i)
+					pp.addvars(nodetbl[node][4])
-				obj.write("usemtl %s\n" % nname.replace(":", "__"))
+					pp.setvar("TEXTURES", "1" if getarg(optargs, "-t") == "" else "0")
-				objd = open("models/" + nodetbl[nname][0] + ".obj", 'r')
+					for line in objdef:
-				pp = Preprocessor(omit_empty=True)
+						line = pp.process(line.rstrip("\r\n"))
-				pp.addvars(parse_arglist(nodetbl[nname][4]))
+						if line is None:
-				pp.setvar("TEXTURES", str(('-t', '') in optargs))
+							continue
-				for line in objd:
+						# Translate vertice coordinates
-					line = pp.process(line[:-1])
+						if line.startswith("v "):
-					if line is None:
+							vx, vy, vz = (float(e) for e in line.split(" ")[1:])
-						continue
+							vx += x
-					if line.startswith("v "):
+							vy += y
-						tmp = line.split(" ")
+							vz += z
-						vx, vy, vz = float(tmp[1]), float(tmp[2]), float(tmp[3])
+							obj.write("v %.1f %.1f %.1f\n" % (vx, vy, vz))
-						vx += x
+						else:
-						vy += y
+							obj.write(line)
-						vz += z
+							obj.write("\n")
-						obj.write("v %.1f %.1f %.1f\n" % (vx, vy, vz))
-					else:
-						obj.write(line + "\n")
-				del pp
-				objd.close()
 				obj.write("\n")
 				i += 1
-	obj.close()
+
-	mtl = open(filepart + ".mtl", "w")
+with open(mtlfile, "w") as mtl:
-	mtl.write("# Generated by mt2obj\n\n")
+	mtl.write("# Generated by mt2obj (https://github.com/sfan5/mt2obj)\n\n")
-	for node in foundnodes:
+	for node in nodes_seen:
+		if node not in nodetbl.keys():
+			continue
 		mtl.write("newmtl %s\n" % node.replace(":", "__"))
 		c = nodetbl[node]
-		mtld = open("models/" + nodetbl[node][0] + ".mtl", 'r')
+		with open("models/%s.mtl" % c[0], "r") as mtldef:
-		pp = Preprocessor(omit_empty=True)
+			pp = Preprocessor(omit_empty=True)
-		pp.addvars(parse_arglist(nodetbl[node][4]))
+			pp.addvars(c[4])
-		pp.addvars({
+			pp.addvars({
-			'r': str(c[1]/255),
+				"r": str(c[1]/255),
-			'g': str(c[2]/255),
+				"g": str(c[2]/255),
-			'b': str(c[3]/255),
+				"b": str(c[3]/255),
-			'a': str(c[5]/255 if len(c) > 5 else 1.0),
+				"a": str(c[5]/255 if len(c) > 5 else 1.0),
-			'TEXTURES': str(('-t', '') in optargs),
+				"TEXTURES": "1" if getarg(optargs, "-t") == "" else "0",
-		})
+			})
-		for line in mtld:
+			for line in mtldef:
-			line = pp.process(line[:-1])
+				line = pp.process(line.rstrip("\r\n"))
-			if line is None:
+				if line is None:
-				continue
+					continue
-			mtl.write(line + "\n")
+				mtl.write(line + "\n")
-		del pp
 		mtl.write("\n")
-		mtld.close()
+
-	mtl.close()
+nodes_unknown = nodes_seen - set(nodetbl.keys())
-	if len(unknownnodes) > 0:
+if len(nodes_unknown) > 0:
-		print("There were some unknown nodes that were ignored during conversion:")
+	print("There were some unknown nodes that were ignored during conversion:")
-		for e in unknownnodes:
+	for node in nodes_unknown:
-			print(e)
+		print("  " + node)
+