summaryrefslogtreecommitdiff
path: root/structutil.py
blob: 7a26e86589c5c1ac08be743ed57d085a9c86bde0 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
import struct

ENDIANESS = '<' # small endian

uint16 = 'H'
uint32 = 'I'
uint64 = 'Q'
uint8 = 'B'

uint32_arr = lambda s: ENDIANESS + uint32*s.e.size

def unpack_list(structs, s):
    """ struct has a bug - 
    sizeof('IBIII') == 20
    sizeof('IIIIB') == 16
    """
    ret = []
    start = 0
    for the_struct in structs:
        ret.extend(list(the_struct.unpack(s[start:start + s.size])))
        t += s.size
    return ret

def group(format):
    return reduce(lambda cs, c: cs[:-1]+[cs[-1]+c] if len(cs) > 0 and c == cs[-1][-1] else cs+[c], format, [])

class Elements(object):
    pass

class StructList(object):

    def __init__(self, formats):
        self._s = map(struct.Struct, (ENDIANESS+f for f in formats))
        self.size = sum([s.size for s in self._s])

    def pack(self, *args):
        i_s, i_e = 0, 0
        r = []
        for s in self._s:
            i_e += len(s.format) - (1 if s.format[0] in '<>' else 0)
            r.append(s.pack(*args[i_s:i_e]))
            i_s = i_e
        return ''.join(r)

    def unpack(self, st):
        r = []
        i_s, i_e = 0, 0
        for s in self._s:
            i_e += s.size
            r.append(list(s.unpack(st[i_s:i_e])))
            i_s = i_e
        return sum(r, [])

class StructMeta(type):
    def __new__(meta, classname, bases, classDict):
        fields = classDict['fields']
        is_complex = classDict['_is_complex'] = callable(fields[-1][0])
        if is_complex:
            classDict['complex_field'] = complex_field = fields[-1]
            fields = fields[:-1]
        assert(not any(map(callable, fields)))
        classDict['_s'] = StructList(group(''.join(t for t,n in fields)))
        classDict['_names'] = [n for t,n in fields]
        classDict['size'] = classDict['_s'].size
        classDict['field_elements'] = [len(t) for t, n in fields]
        return type.__new__(meta, classname, bases, classDict)

def slice_pairs_iter(sizes):
    s = 0
    for size in sizes:
        yield s, s+size
        s += size

def cut(elements, sizes):
    for s, e in slice_pairs_iter(sizes):
        if e - s == 1:
            yield elements[s]
        else:
            yield elements[s:e]

class Struct(object):
    @classmethod
    def parse(cls, s):
        base = list(cut(cls._s.unpack(s[:cls._s.size]), cls.field_elements))
        if cls._is_complex:
            import pdb; pdb.set_trace()
            return base
        return base

    @classmethod
    def make(cls, **kw):
        args = []
        args = [kw[n] for n in cls._names]
        assert(len(args) == len(cls._names) == len(kw))
        return cls._s.pack(*args)

    def __init__(self, *args, **kw):
        self.e = Elements()
        if (len(kw) == 0 and len(args) == 1) or (len(kw) == 1 and kw.has_key('s')):
            s = args[0] if len(args) == 1 else kw['s']
            self.elements = elements = self.parse(s)
        else:
            self.elements = elements = [kw[n] for n in self._names]
        self.e.__dict__.update(dict(zip(self._names, elements)))

    def tostr(self):
        return self.make(**self.e.__dict__)

def list_to_str(l, type=uint32):
    if len(l) == 0:
        return ''
    return struct.pack(ENDIANESS+len(l)*type, l)