# Ideally we'd probably be using something like the PyObjC runtime
# rather than the struct module to get by here. However, this is
# good enough for now.
import struct
from itertools import *
__all__ = """
sizeof
Structure
pypackable
p_char
p_byte
p_ubyte
p_short
p_ushort
p_int
p_uint
p_long
p_ulong
p_float
p_double
p_ptr
p_longlong
p_ulonglong
""".split()
def sizeof(s):
"""
Return the size of an object when packed
"""
if hasattr(s, '_size_'):
return s._size_
return len(s)
class MetaPackable(type):
"""
Fixed size struct.unpack-able types use from_tuple as their designated initializer
"""
def from_mmap(cls, mm, ptr, **kw):
return cls.from_str(mm[ptr:ptr+cls._size_], **kw)
def from_fileobj(cls, f, **kw):
return cls.from_str(f.read(cls._size_), **kw)
def from_str(cls, s, **kw):
endian = kw.get('_endian_', cls._endian_)
return cls.from_tuple(struct.unpack(endian + cls._format_, s), **kw)
def from_tuple(cls, tpl, **kw):
return cls(tpl[0], **kw)
class BasePackable(object):
# XXX - use big endian everywhere, because we're only parsing Mach-O
_endian_ = '>'
def to_str(self, s):
raise NotImplementedError
def to_fileobj(self, f):
f.write(self.to_str())
def to_mmap(self, mm, ptr):
mm[ptr:ptr+self._size_] = self.to_str()
class Packable(BasePackable):
"""
Fixed size single object that is (un)packable by a struct format
"""
__metaclass__ = MetaPackable
def to_str(self):
cls = type(self)
endian = getattr(self, '_endian_', cls._endian_)
return struct.pack(endian + cls._format_, self)
def pypackable(name, pytype, format):
"""
Create a "mix-in" class with a python type and a
Packable with the given struct format
"""
size, items = formatinfo(format)
return type(Packable)(name, (pytype, Packable), {
'_format_': format,
'_size_': size,
'_items_': items,
})
def formatinfo(format):
"""
Calculate the size and number of items in a struct format.
"""
size = struct.calcsize(format)
return size, len(struct.unpack(format, '\x00' * size))
class MetaStructure(MetaPackable):
"""
The metaclass of Structure objects that does all the magic.
Since we can assume that all Structures have a fixed size,
we can do a bunch of calculations up front and pack or
unpack the whole thing in one struct call.
"""
def __new__(cls, clsname, bases, dct):
fields = dct['_fields_']
names = []
types = []
structmarks = []
format = ''
items = 0
size = 0
def struct_property(name, typ):
def _get(self):
return self._objects_[name]
def _set(self, obj):
if type(obj) is not typ:
obj = typ(obj)
self._objects_[name] = obj
return property(_get, _set, typ.__name__)
for name, typ in fields:
dct[name] = struct_property(name, typ)
names.append(name)
types.append(typ)
format += typ._format_
size += typ._size_
if (typ._items_ > 1):
structmarks.append((items, typ._items_, typ))
items += typ._items_
dct['_structmarks_'] = structmarks
dct['_names_'] = names
dct['_types_'] = types
dct['_size_'] = size
dct['_items_'] = items
dct['_format_'] = format
return super(MetaStructure, cls).__new__(cls, clsname, bases, dct)
def from_tuple(cls, tpl, **kw):
values = []
current = 0
for begin, length, typ in cls._structmarks_:
if begin > current:
values.extend(tpl[current:begin])
current = begin + length
values.append(typ.from_tuple(tpl[begin:current], **kw))
values.extend(tpl[current:])
return cls(*values, **kw)
class Structure(BasePackable):
__metaclass__ = MetaStructure
_fields_ = ()
def __init__(self, *args, **kwargs):
if len(args) == 1 and not kwargs and type(args[0]) is type(self):
kwargs = args[0]._objects_
args = ()
self._objects_ = {}
iargs = chain(izip(self._names_, args), kwargs.iteritems())
for key, value in iargs:
if key not in self._names_ and key != "_endian_":
raise TypeError
setattr(self, key, value)
for key, typ in izip(self._names_, self._types_):
if key not in self._objects_:
self._objects_[key] = typ()
def _get_packables(self):
for obj in imap(self._objects_.__getitem__, self._names_):
if obj._items_ == 1:
yield obj
else:
for obj in obj._get_packables():
yield obj
def to_str(self):
return struct.pack(self._endian_ + self._format_, *self._get_packables())
def __cmp__(self, other):
if not type(other) is type(self):
raise TypeError, 'Cannot compare objects of type %r to objects of type %r' % (type(other), type(self))
for cmpval in starmap(cmp, izip(self._get_packables(), other._get_packables())):
if cmpval != 0:
return cmpval
return 0
# export common packables with predictable names
p_char = pypackable('p_char', str, 'c')
p_byte = pypackable('p_byte', int, 'b')
p_ubyte = pypackable('p_ubyte', int, 'B')
p_short = pypackable('p_short', int, 'h')
p_ushort = pypackable('p_ushort', int, 'H')
p_int = pypackable('p_int', int, 'i')
p_uint = pypackable('p_uint', long, 'I')
p_long = pypackable('p_long', int, 'l')
p_ulong = pypackable('p_ulong', long, 'L')
p_float = pypackable('p_float', float, 'f')
p_double = pypackable('p_double', float, 'd')
p_ptr = pypackable('p_ptr', long, 'P')
p_longlong = pypackable('p_longlong', long, 'q')
p_ulonglong = pypackable('p_ulonglong', long, 'Q')
def test_ptypes():
from cStringIO import StringIO
import mmap
class MyStructure(Structure):
_fields_ = (
('foo', p_int),
('bar', p_ubyte),
)
class MyFunStructure(Structure):
_fields_ = (
('fun', p_char),
('mystruct', MyStructure),
)
for endian in '><':
kw = dict(_endian_=endian)
MYSTRUCTURE = '\x00\x11\x22\x33\xFF'
for fn, args in [
('from_str', (MYSTRUCTURE,)),
('from_mmap', (MYSTRUCTURE, 0)),
('from_fileobj', (StringIO(MYSTRUCTURE),)),
]:
myStructure = getattr(MyStructure, fn)(*args, **kw)
if endian == '>':
assert myStructure.foo == 0x00112233
else:
assert myStructure.foo == 0x33221100
assert myStructure.bar == 0xFF
assert myStructure.to_str() == MYSTRUCTURE
MYFUNSTRUCTURE = '!' + MYSTRUCTURE
for fn, args in [
('from_str', (MYFUNSTRUCTURE,)),
('from_mmap', (MYFUNSTRUCTURE, 0)),
('from_fileobj', (StringIO(MYFUNSTRUCTURE),)),
]:
myFunStructure = getattr(MyFunStructure, fn)(*args, **kw)
assert myFunStructure.mystruct == myStructure
assert myFunStructure.fun == '!'
assert myFunStructure.to_str() == MYFUNSTRUCTURE
sio = StringIO()
myFunStructure.to_fileobj(sio)
assert sio.getvalue() == MYFUNSTRUCTURE
mm = mmap.mmap(-1, sizeof(MyFunStructure) * 2, mmap.MAP_ANONYMOUS)
mm[:] = '\x00' * (sizeof(MyFunStructure) * 2)
myFunStructure.to_mmap(mm, 0)
assert MyFunStructure.from_mmap(mm, 0, **kw) == myFunStructure
assert mm[:sizeof(MyFunStructure)] == MYFUNSTRUCTURE
assert mm[sizeof(MyFunStructure):] == '\x00' * sizeof(MyFunStructure)
myFunStructure.to_mmap(mm, sizeof(MyFunStructure))
assert mm[:] == MYFUNSTRUCTURE + MYFUNSTRUCTURE
assert MyFunStructure.from_mmap(mm, sizeof(MyFunStructure), **kw) == myFunStructure
if __name__ == '__main__':
test_ptypes()
|