structs.py

import bitstring


class StructValueError(Exception): pass


class Member(object):
	creation_counter = 0
	abstract = False

	def __init__(self, format, default=None, pos=None):
		self.creation_counter = Member.creation_counter
		Member.creation_counter += 1

		self.format = format
		self.default = default
		self.pos = pos
		self.name = None

	@property
	def length(self):
		stretchy, tokens = bitstring.tokenparser(self.format)
		if len(tokens) != 1:
			raise ValueError("Format string must contain exactly one argument")
		if stretchy or tokens[0][1] is None:
			raise ValueError("Only fixed-length format arguments are supported.")
		return tokens[0][1]

	def __get__(self, obj, type=None):
		if type is None:
			return self
		return getattr(obj, '_' + self.name, self.default)

	def __set__(self, obj, value):
		setattr(obj, '_' + self.name, value)

	def frombits(self, obj, bits):
		bits.pos = self.pos
		self.__set__(obj, bits.read(self.format))

	def tobits(self, obj, bits):
		bits[self.pos:self.pos+self.length] = bitstring.pack(self.format, self.__get__(obj, type(obj)))


class AbstractMember(Member):
	abstract = True


class ConstantMember(Member):
	def __init__(self, format, value, *args, **kwargs):
		super(ConstantMember, self).__init__(format, *args, **kwargs)
		self.value = value

	def __set__(self, obj, value):
		if value != self.value:
			raise StructValueError()

	def __get__(self, obj, type=None):
		if type is not None:
			return self
		return self.value


class StructMeta(type):
	def __new__(self, name, bases, attrs):
		klass = super(StructMeta, self).__new__(name, bases, attrs)
		fields = []
		for base in bases:
			if hasattr(base, '_fields'):
				fields.extend(base._fields)
		pos = 0
		abstract = False
		for name, field in attrs.items():
			if isinstance(field, Member):
				field.name = name

				if field.pos is not None:
					pos = field.pos
				else:
					field.pos = pos
				pos += field.length

				fields.append((name, field))

				abstract = abstract or field.abstract
		fields.sort(key=lambda field: field.creation_counter)
		klass._fields = fields
		klass._abstract = abstract


class Struct(object):
	__metaclass__ = StructMeta

	def __init__(self, *args, **kwargs):
		for i, arg in enumerate(args):
			setattr(self, self._fields[i], arg)
		for name, value in kwargs.items():
			setattr(self, name, value)

	@classmethod
	def decode(cls, data):
		bits = bitstring.BitString(bytes=data)
		ret = cls()
		for field in cls._fields:
			field.frombits(ret, bits)