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)"