1 | """StarStruct element class.""" |
||
2 | |||
3 | 1 | import struct |
|
4 | 1 | import re |
|
5 | |||
6 | 1 | from starstruct.element import register, Element |
|
7 | 1 | from starstruct.modes import Mode |
|
8 | |||
9 | |||
10 | 1 | @register |
|
11 | 1 | class ElementLength(Element): |
|
12 | """ |
||
13 | The length StarStruct element class. |
||
14 | """ |
||
15 | |||
16 | 1 | View Code Duplication | def __init__(self, field, mode=Mode.Native, alignment=1): |
0 ignored issues
–
show
Duplication
introduced
by
![]() |
|||
17 | """Initialize a StarStruct element object.""" |
||
18 | |||
19 | # All of the type checks have already been performed by the class |
||
20 | # factory |
||
21 | 1 | if isinstance(field[0], str): |
|
22 | 1 | self.name = field[0] |
|
23 | 1 | self.object_length = True |
|
24 | 1 | elif isinstance(field[0], bytes): |
|
25 | 1 | self.name = field[0].decode('utf-8') |
|
26 | 1 | self.object_length = False |
|
27 | |||
28 | 1 | self.ref = field[2] |
|
29 | |||
30 | 1 | self._mode = mode |
|
31 | 1 | self._alignment = alignment |
|
32 | |||
33 | # Validate that the format specifiers are valid struct formats, this |
||
34 | # doesn't have to be done now because the format will be checked when |
||
35 | # any struct functions are called, but it's better to inform the user of |
||
36 | # any errors earlier. |
||
37 | # The easiest way to perform this check is to create a "Struct" class |
||
38 | # instance, this will also increase the efficiency of all struct related |
||
39 | # functions called. |
||
40 | 1 | self.format = mode.value + field[1] |
|
41 | 1 | self._struct = struct.Struct(self.format) |
|
42 | |||
43 | 1 | @staticmethod |
|
44 | def valid(field): |
||
45 | """ |
||
46 | Validation function to determine if a field tuple represents a valid |
||
47 | enum element type. |
||
48 | |||
49 | The basics have already been validated by the Element factory class, |
||
50 | validate that the struct format is a valid unsigned numeric value. |
||
51 | """ |
||
52 | 1 | return len(field) == 3 \ |
|
53 | and isinstance(field[1], str) \ |
||
54 | and re.match(r'[BHILQ]', field[1]) \ |
||
55 | and isinstance(field[2], str) and len(field[2]) |
||
56 | |||
57 | 1 | def validate(self, msg): |
|
58 | """ |
||
59 | Ensure that the supplied message contains the required information for |
||
60 | this element object to operate. |
||
61 | |||
62 | All elements that are Variable must reference valid Length elements. |
||
63 | """ |
||
64 | # TODO: Allow referencing multiple elements for byte lengths? |
||
65 | |||
66 | 1 | from starstruct.elementvariable import ElementVariable |
|
67 | 1 | if not isinstance(msg[self.ref], ElementVariable): |
|
68 | err = 'length field {} reference {} invalid type' |
||
69 | raise TypeError(err.format(self.name, self.ref)) |
||
70 | 1 | elif not msg[self.ref].ref == self.name: |
|
71 | err = 'length field {} reference {} mismatch' |
||
72 | raise TypeError(err.format(self.name, self.ref)) |
||
73 | |||
74 | 1 | def update(self, mode=None, alignment=None): |
|
75 | """change the mode of the struct format""" |
||
76 | 1 | if alignment: |
|
77 | 1 | self._alignment = alignment |
|
78 | |||
79 | 1 | if mode: |
|
80 | 1 | self._mode = mode |
|
81 | 1 | self.format = mode.value + self.format[1:] |
|
82 | # recreate the struct with the new format |
||
83 | 1 | self._struct = struct.Struct(self.format) |
|
84 | |||
85 | 1 | def pack(self, msg): |
|
86 | """Pack the provided values into the supplied buffer.""" |
||
87 | 1 | if self.object_length: |
|
88 | # When packing a length element, use the length of the referenced |
||
89 | # element not the value of the current element in the supplied |
||
90 | # object. |
||
91 | 1 | data = self._struct.pack(len(msg[self.ref])) |
|
92 | else: |
||
93 | # When packing something via byte length, |
||
94 | # we use our self to determine the length |
||
95 | 1 | data = self._struct.pack(msg[self.name]) |
|
96 | |||
97 | # If the data does not meet the alignment, add some padding |
||
98 | 1 | missing_bytes = len(data) % self._alignment |
|
99 | 1 | if missing_bytes: |
|
100 | data += b'\x00' * missing_bytes |
||
101 | 1 | return data |
|
102 | |||
103 | 1 | def unpack(self, msg, buf): |
|
104 | """Unpack data from the supplied buffer using the initialized format.""" |
||
105 | 1 | ret = self._struct.unpack_from(buf, 0) |
|
106 | |||
107 | # Remember to remove any alignment-based padding |
||
108 | 1 | extra_bytes = self._alignment - 1 - (struct.calcsize(self.format) % |
|
109 | self._alignment) |
||
110 | 1 | unused = buf[struct.calcsize(self.format) + extra_bytes:] |
|
111 | 1 | return (ret[0], unused) |
|
112 | |||
113 | 1 | def make(self, msg): |
|
114 | """Return the length of the referenced array""" |
||
115 | 1 | if self.object_length: |
|
116 | 1 | return len(msg[self.ref]) |
|
117 | else: |
||
118 | return msg[self.name] |
||
119 |