binstruct.py

from __future__ import print_function
import struct
from omg.util import *

class StructMeta(type):

	@staticmethod
	def _get(f):
		# Get struct field or default value
		def fn(self):
			try:
				return self._values[f[0]]
			except KeyError:
				self._values[f[0]] = f[2]
				return f[2]
		return fn
		
	@staticmethod
	def _set(f):
		# Set struct field, stripping null bytes from string fields
		def fn(self, value):
			if 's' in f[1]:
				self._values[f[0]] = safe_name(zstrip(value))
			else:
				self._values[f[0]] = value
		return fn

	@property
	def size(cls):
		return cls._struct.size

	def __new__(cls, name, bases, dict):
	
		fields = dict.get("__fields__", [])
		
		# Set up attributes for all defined fields
		for f in fields:
			field_name = f[0]
			field_doc  = f[3] if len(f) > 3 else None
			
			dict[field_name] = property(StructMeta._get(f), StructMeta._set(f), doc = field_doc)
		
		# TODO: set up optional bitfields also (for linedef flags, etc)
		
		# Set up struct format string and size
		dict["_keys"] = [f[0] for f in fields if f[1] != 'x']
		dict["_struct"] = struct.Struct("<" + "".join(f[1] for f in fields))
		
		return type.__new__(cls, name, bases, dict)

# Python 2 and 3 metaclass hack
_StructParent = StructMeta("_StructParent", (object,), {})

class Struct(_StructParent):
	"""
	Class which allows easily creating additional classes based on binary structs.
	
	Create a subclass of Struct with an attribute called __fields__.
	This is a list consisting of tuples with the following info:
	 - name: the name of a struct field
	 - type: the type of the field, such as 'h' or '8s' (see Python's 'struct' module)
	 - value: default value to use
	 - docstring (optional)
	
	Class properties will be automatically generated for all struct members.
	Strings (type 's') will be padded with null bytes on packing, or stripped on unpacking.
	Since this is Doom WAD stuff, some string sanitization will also happen (and they'll
	always be padded to 8 chars since it's assumed to be a lump name).
	
	The 'pack' and 'unpack' methods convert a struct instance to/from binary data.
	"""
	
	def __init__(self, *args, **kwargs):
		"""
		Create a new instance of this struct.
		Arguments can be either positional (based on the layout of the struct),
		and/or keyword arguments based on the names of the struct members.
		
		Other optional arguments:
		'bytes' - a bytestring. The struct will be unpacked from this data
			(other args will be ignored.)
		"""
		self._values = {}
		if 'bytes' in kwargs:
			self.unpack(kwargs['bytes'])
		else:
			values = {}
			values.update(dict(zip(self._keys, args)))
			values.update(kwargs)
			for key, value in values.items():
				if key in self._keys:
					setattr(self, key, value)
	
	def pack(self):
		packs = []
		for f in self.__fields__:
			if 's' in f[1]:
				packs.append(zpad(safe_name(getattr(self, f[0]))))
			elif f[1] != 'x':
				packs.append(getattr(self, f[0]))
		return self._struct.pack(*packs)

	def unpack(self, data):
		for key, value in zip(self._keys, self._struct.unpack(data)):
			setattr(self, key, value)

class Vertex(Struct):
	"""Represents a map vertex"""
	
	__fields__ = [
	("x", "h", 0),
	("y", "h", 0),
	]
	
	def foo(self):
		print("My coordinates are ({0}, {1}).".format(self.x, self.y))

if __name__ == "__main__":
	v1 = Vertex()
	v1.foo() # "My coordinates are (0, 0)"
	
	v2 = Vertex(0, 0)
	v2.foo() # "My coordinates are (0, 0)"
	
	v3 = Vertex(y = 10)
	v3.foo() # "My coordinates are (0, 10)"
	
	v4 = Vertex(bytes = b"\x05\x00\x0a\x00")
	v4.foo() # "My coordinates are (5, 10)"