Issues (43)

starstruct/elementstring.py (1 issue)

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 ElementString(Element):
12
    """
13
    A StarStruct element for strings, because standard string treatment of
14
    pack/unpack can be inconvenient.
15
16
    This element will encode and decode string type elements from and to forms
17
    that are easier to use and manage.
18
    """
19
20 1
    def __init__(self, field, mode=Mode.Native, alignment=1):
21
        """Initialize a StarStruct element object."""
22
23
        # All of the type checks have already been performed by the class
24
        # factory
25 1
        self.name = field[0]
26
27
        # The ref attribute is required for all elements, but base element
28
        # types don't have one
29 1
        self.ref = None
30
31 1
        self._mode = mode
32 1
        self._alignment = alignment
33
34
        # Validate that the format specifiers are valid struct formats, this
35
        # doesn't have to be done now because the format will be checked when
36
        # any struct functions are called, but it's better to inform the user of
37
        # any errors earlier.
38
        # The easiest way to perform this check is to create a "Struct" class
39
        # instance, this will also increase the efficiency of all struct related
40
        # functions called.
41 1
        self.format = mode.value + field[1]
42 1
        self._struct = struct.Struct(self.format)
43
44 1
    @staticmethod
45
    def valid(field):
46
        """
47
        Validation function to determine if a field tuple represents a valid
48
        string element type.
49
        """
50 1
        return len(field) == 2 \
51
            and isinstance(field[1], str) \
52
            and re.match(r'\d*[csp]', field[1])
53
54 1
    def validate(self, msg):
55
        """
56
        Ensure that the supplied message contains the required information for
57
        this element object to operate.
58
59
        The "string" element requires no further validation.
60
        """
61 1
        pass
62
63 1
    def update(self, mode=None, alignment=None):
64
        """change the mode of the struct format"""
65 1
        if alignment:
66 1
            self._alignment = alignment
67
68 1
        if mode:
69 1
            self._mode = mode
70 1
            self.format = mode.value + self.format[1:]
71
            # recreate the struct with the new format
72 1
            self._struct = struct.Struct(self.format)
73
74 1
    def pack(self, msg):
75
        """Pack the provided values into the supplied buffer."""
76
        # Ensure that the input is of the proper form to be packed
77 1
        val = msg[self.name]
78 1
        size = struct.calcsize(self.format)
79 1
        assert len(val) <= size
80 1
        if self.format[-1] in ('s', 'p'):
81 1
            if not isinstance(val, bytes):
82 1
                assert isinstance(val, str)
83 1
                val = val.encode()
84 1
                if self.format[-1] == 'p' and len(val) < size:
85
                    # 'p' (pascal strings) must be the exact size of the format
86 1
                    val += b'\x00' * (size - len(val))
87
88 1
            data = self._struct.pack(val)
89
        else:  # 'c'
90 1
            if not all(isinstance(c, bytes) for c in val):
91 1
                if isinstance(val, bytes):
92
                    val = [bytes([c]) for c in val]
93
                else:
94
                    # last option, it could be a string, or a list of strings
95 1
                    assert (isinstance(val, list) and
96
                            all(isinstance(c, str) for c in val)) or \
97
                        isinstance(val, str)
98 1
                    val = [c.encode() for c in val]
99 1
            if len(val) < size:
100 1
                val.extend([b'\x00'] * (size - len(val)))
101 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...
102
103
        # If the data does not meet the alignment, add some padding
104 1
        missing_bytes = len(data) % self._alignment
105 1
        if missing_bytes:
106 1
            data += b'\x00' * (self._alignment - missing_bytes)
107 1
        return data
108
109 1
    def unpack(self, msg, buf):
110
        """Unpack data from the supplied buffer using the initialized format."""
111 1
        ret = self._struct.unpack_from(buf, 0)
112
113
        # Remember to remove any alignment-based padding
114 1
        extra_bytes = self._alignment - 1 - (struct.calcsize(self.format) %
115
                                             self._alignment)
116 1
        unused = buf[struct.calcsize(self.format) + extra_bytes:]
117
118 1
        if self.format[-1] in 's':
119
            # for 's' formats, convert to a string and strip padding
120 1
            val = ret[0].decode().strip('\x00')
121 1
        elif self.format[-1] in 'p':
122
            # for 'p' formats, convert to a string, but leave the padding
123 1
            val = ret[0].decode()
124
        else:  # 'c'
125
            # Just in case we have some ints in the message
126 1
            val = [c.decode() if not isinstance(c, int)
127
                   else chr(c)
128
                   for c in ret]
129 1
        return (val, unused)
130
131 1
    def make(self, msg):
132
        """Return a string of the expected format"""
133 1
        val = msg[self.name]
134 1
        size = struct.calcsize(self.format)
135 1
        assert len(val) <= size
136
137
        # If the supplied value is a list of chars, or a list of bytes, turn
138
        # it into a string for ease of processing.
139 1
        if isinstance(val, list):
140 1
            if all(isinstance(c, bytes) for c in val):
141
                val = ''.join([c.decode() for c in val])
142 1
            elif all(isinstance(c, str) for c in val):
143 1
                val = ''.join([c for c in val])
144
            else:
145 1
                error = 'Invalid value for string element: {}'
146 1
                raise TypeError(error.format(val))
147 1
        elif isinstance(val, bytes):
148
            # If the supplied value is a byes, decode it into a normal string
149
            val = val.decode()
150
151
        # 'p' (pascal strings) and 'c' (char list) must be the exact size of
152
        # the format
153 1
        if self.format[-1] == 'p' and len(val) < size - 1:
154 1
            val += '\x00' * (size - len(val) - 1)
155
156
        # Lastly, 'c' (char list) formats are expected to be a list of
157
        # characters rather than a string.
158 1
        if self.format[-1] == 'c':
159 1
            val += '\x00' * (size - len(val))
160 1
            val = [c for c in val]
161
162
        return val
163