Skip to content

Instantly share code, notes, and snippets.

@devinacker
Last active March 7, 2021 23:52
Show Gist options
  • Save devinacker/65cb931709d601d832e43852d922abe0 to your computer and use it in GitHub Desktop.
Save devinacker/65cb931709d601d832e43852d922abe0 to your computer and use it in GitHub Desktop.
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)"
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment