1 | """StarStruct element class.""" |
||
2 | |||
3 | 1 | import struct |
|
4 | 1 | import re |
|
5 | 1 | import enum |
|
6 | |||
7 | 1 | from starstruct.element import register, Element |
|
8 | 1 | from starstruct.modes import Mode |
|
9 | |||
10 | |||
11 | 1 | @register |
|
12 | 1 | class ElementNum(Element): |
|
13 | """ |
||
14 | A StarStruct element class for number fields. |
||
15 | """ |
||
16 | # pylint: disable=too-many-instance-attributes |
||
17 | |||
18 | 1 | def __init__(self, field, mode=Mode.Native, alignment=1): |
|
19 | """Initialize a StarStruct element object.""" |
||
20 | |||
21 | # All of the type checks have already been performed by the class |
||
22 | # factory |
||
23 | 1 | self.name = field[0] |
|
24 | |||
25 | # The ref attribute is required for all elements, but the base element |
||
26 | # type does not have one |
||
27 | 1 | self.ref = None |
|
28 | |||
29 | 1 | self._mode = mode |
|
30 | 1 | self._alignment = alignment |
|
31 | |||
32 | # Validate that the format specifiers are valid struct formats, this |
||
33 | # doesn't have to be done now because the format will be checked when |
||
34 | # any struct functions are called, but it's better to inform the user of |
||
35 | # any errors earlier. |
||
36 | # The easiest way to perform this check is to create a "Struct" class |
||
37 | # instance, this will also increase the efficiency of all struct related |
||
38 | # functions called. |
||
39 | 1 | self.format = mode.value + field[1] |
|
40 | 1 | self._struct = struct.Struct(self.format) |
|
41 | |||
42 | # for numeric elements we should also keep track of how many numeric |
||
43 | # fields and what the size of those fields are required to create this |
||
44 | # element. |
||
45 | 1 | self._bytes = struct.calcsize(self.format[-1]) |
|
46 | 1 | self._signed = self.format[-1] in 'bhilq' |
|
47 | |||
48 | 1 | @staticmethod |
|
49 | def valid(field): |
||
50 | """ |
||
51 | Validation function to determine if a field tuple represents a valid |
||
52 | base element type. |
||
53 | |||
54 | The basics have already been validated by the Element factory class, |
||
55 | validate the specific struct format now. |
||
56 | """ |
||
57 | 1 | return len(field) == 2 \ |
|
58 | and isinstance(field[1], str) \ |
||
59 | and re.match(r'\d*[bBhHiIlLqQ]', field[1]) |
||
60 | |||
61 | 1 | def validate(self, msg): |
|
62 | """ |
||
63 | Ensure that the supplied message contains the required information for |
||
64 | this element object to operate. |
||
65 | |||
66 | The "number" element requires no further validation. |
||
67 | """ |
||
68 | 1 | pass |
|
69 | |||
70 | 1 | def update(self, mode=None, alignment=None): |
|
71 | """change the mode of the struct format""" |
||
72 | 1 | if alignment: |
|
73 | 1 | self._alignment = alignment |
|
74 | |||
75 | 1 | if mode: |
|
76 | 1 | self._mode = mode |
|
77 | 1 | self.format = mode.value + self.format[1:] |
|
78 | # recreate the struct with the new format |
||
79 | 1 | self._struct = struct.Struct(self.format) |
|
80 | |||
81 | 1 | def pack(self, msg): |
|
82 | """Pack the provided values into the supplied buffer.""" |
||
83 | # Take a single numeric value and convert it into the necessary list |
||
84 | # of values required by the specified format. |
||
85 | 1 | val = msg[self.name] |
|
86 | |||
87 | # This should be a number, but handle cases where it's an enum |
||
88 | 1 | if isinstance(val, enum.Enum): |
|
89 | val = val.value |
||
90 | |||
91 | # If the value supplied is not already a bytes object, convert it now. |
||
92 | 1 | if isinstance(val, (bytes, bytearray)): |
|
93 | val_list = val |
||
94 | else: |
||
95 | 1 | val_list = val.to_bytes(struct.calcsize(self.format), |
|
96 | byteorder=self._mode.to_byteorder(), |
||
97 | signed=self._signed) |
||
98 | |||
99 | # join the byte list into the expected number of values to pack the |
||
100 | # specified struct format. |
||
101 | 1 | val = [int.from_bytes(val_list[i:i + self._bytes], # pylint: disable=no-member |
|
102 | byteorder=self._mode.to_byteorder(), |
||
103 | signed=self._signed) |
||
104 | for i in range(0, len(val_list), self._bytes)] |
||
105 | 1 | data = self._struct.pack(*val) |
|
0 ignored issues
–
show
|
|||
106 | |||
107 | # If the data does not meet the alignment, add some padding |
||
108 | 1 | missing_bytes = len(data) % self._alignment |
|
109 | 1 | if missing_bytes: |
|
110 | data += b'\x00' * missing_bytes |
||
111 | 1 | return data |
|
112 | |||
113 | 1 | def unpack(self, msg, buf): |
|
114 | """Unpack data from the supplied buffer using the initialized format.""" |
||
115 | 1 | ret = self._struct.unpack_from(buf, 0) |
|
116 | |||
117 | # Remember to remove any alignment-based padding |
||
118 | 1 | extra_bytes = self._alignment - 1 - (struct.calcsize(self.format) % |
|
119 | self._alignment) |
||
120 | 1 | unused = buf[struct.calcsize(self.format) + extra_bytes:] |
|
121 | |||
122 | # merge the unpacked data into a byte array |
||
123 | 1 | data = [v.to_bytes(self._bytes, byteorder=self._mode.to_byteorder(), |
|
124 | signed=self._signed) for v in ret] |
||
125 | # Join the returned list of numbers into a single value |
||
126 | 1 | val = int.from_bytes(b''.join(data), # pylint: disable=no-member |
|
127 | byteorder=self._mode.to_byteorder(), |
||
128 | signed=self._signed) |
||
129 | 1 | return (val, unused) |
|
130 | |||
131 | 1 | def make(self, msg): |
|
132 | """Return the expected "made" value""" |
||
133 | 1 | val = msg[self.name] |
|
134 | |||
135 | # This should be a number, but handle cases where it's an enum |
||
136 | 1 | if isinstance(val, enum.Enum): |
|
137 | val = val.value |
||
138 | 1 | elif isinstance(val, list): |
|
139 | # It's unlikely but possible that this could be a list of numbers, |
||
140 | # or a list of bytes |
||
141 | if all(isinstance(v, bytes) for v in val): |
||
142 | # To turn this into a single number, merge the bytes, later the |
||
143 | # bytes will be converted into a single number. |
||
144 | data = b''.join(val) |
||
145 | elif all(isinstance(v, int) for v in val): |
||
146 | # To turn this into a single number, convert the numbers into |
||
147 | # bytes, and merge the bytes, later the bytes will be converted |
||
148 | # into a single number. |
||
149 | data = [v.to_bytes(self._bytes, |
||
150 | byteorder=self._mode.to_byteorder(), |
||
151 | signed=self._signed) for v in val] |
||
152 | else: |
||
153 | error = 'Invalid value for numerical element: {}' |
||
154 | raise TypeError(error.format(val)) |
||
155 | 1 | elif isinstance(val, bytes): |
|
156 | # If the value supplied is a bytes object, convert it to a number |
||
157 | data = val |
||
158 | 1 | elif isinstance(val, int): |
|
159 | 1 | return val |
|
160 | else: |
||
161 | error = 'Invalid value for numerical element: {}' |
||
162 | raise TypeError(error.format(val)) |
||
163 | |||
164 | return int.from_bytes(data, # pylint: disable=no-member |
||
165 | byteorder=self._mode.to_byteorder(), |
||
166 | signed=self._signed) |
||
167 |
Generally, there is nothing wrong with usage of
*
or**
arguments. For readability of the code base, we suggest to not over-use these language constructs though.For more information, we can recommend this blog post from Ned Batchelder including its comments which also touches this aspect.