Issues (43)

starstruct/elementnum.py (1 issue)

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
Usage of * or ** arguments should usually be done with care.

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.

Loading history...
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