Merge pull request #8 from nocarryr/codspeed

Codspeed

Author: nocarryr

.gitignore

@@ -91,3 +91,4 @@ ENV/
 venv*
 .venv*
 .pytest_cache/
+.codspeed/

pyproject.toml

@@ -59,6 +59,7 @@ test = [
     "faker>=35.2.2",
     "pytest>=8.3.5",
     "pytest-asyncio>=0.24.0",
+    "pytest-codspeed>=4.1.1",
     "pytest-cov>=5.0.0",
     "pytest-doctestplus>=1.3.0",
     "pytest-timeout>=2.4.0",

src/tslumd/messages.py

@@ -103,6 +103,20 @@ class Display:
 
     .. versionadded:: 0.0.2
     """
+    text_length: int|None = field(default=None, compare=False)
+    """Length of the :attr:`text` field
+
+    If provided, the text output will be zero-padded or truncated to this length.
+    If ``None`` (the default), the text length is variable.
+
+
+    .. note::
+
+        This is mainly for internal use during test runs to match the fixed-length
+        text field in actual message bytes.
+
+    .. versionadded:: 0.0.8
+    """
 
     def __post_init__(self):
         self.is_broadcast = self.index == 0xffff
@@ -124,11 +138,23 @@ def broadcast(cls, **kwargs) -> Display:
         return cls(**kwargs)
 
     @classmethod
-    def from_dmsg(cls, flags: Flags, dmsg: bytes) -> Tuple[Display, bytes]:
+    def from_dmsg(cls, flags: Flags, dmsg: bytes, retain_text_length: bool = False) -> Tuple[Display, bytes]:
         """Construct an instance from a ``DMSG`` portion of received message.
 
         Any remaining message data after the relevant ``DMSG`` is returned along
         with the instance.
+
+        Arguments:
+            flags: The message :class:`Flags` field
+            dmsg: The portion of the message containing the ``DMSG`` data
+            retain_text_length: If ``True``, the :attr:`text_length` attribute
+                will be set to the length of the text field as found in the
+                message bytes. Otherwise (the default), it will be set to ``None``
+                and the text length will be variable.
+
+        .. versionchanged:: 0.0.8
+
+            The `retain_text_length` argument was added.
         """
         if len(dmsg) < 4:
             raise DmsgParseError('Invalid dmsg length', dmsg)
@@ -163,9 +189,14 @@ def from_dmsg(cls, flags: Flags, dmsg: bytes) -> Tuple[Display, bytes]:
             if Flags.UTF16 in flags:
                 txt = txt_bytes.decode('UTF-16le')
             else:
+                txt_length = len(txt_bytes)
                 if b'\0' in txt_bytes:
                     txt_bytes = txt_bytes.split(b'\0')[0]
                 txt = txt_bytes.decode('UTF-8')
+                if retain_text_length:
+                    kw['text_length'] = txt_length
+                else:
+                    kw['text_length'] = None
             kw['text'] = txt
         return cls(**kw), dmsg
 
@@ -236,6 +267,8 @@ def to_dmsg(self, flags: Flags) -> bytes:
                 txt_bytes = bytes(self.text, 'UTF-16le')
             else:
                 txt_bytes = bytes(self.text, 'UTF-8')
+            if self.text_length is not None:
+                txt_bytes = txt_bytes.ljust(self.text_length, b'\0')[:self.text_length]
             txt_byte_len = len(txt_bytes)
             data = bytearray(struct.pack('<3H', self.index, ctrl, txt_byte_len))
             data.extend(txt_bytes)
@@ -244,6 +277,7 @@ def to_dmsg(self, flags: Flags) -> bytes:
     def to_dict(self) -> dict:
         d = dataclasses.asdict(self)
         del d['is_broadcast']
+        del d['text_length']
         return d
 
     @classmethod
@@ -352,10 +386,17 @@ def broadcast(cls, **kwargs) -> Message:
         return cls(**kwargs)
 
     @classmethod
-    def parse(cls, msg: bytes) -> Tuple[Message, bytes]:
+    def parse(cls, msg: bytes, retain_text_length: bool = False) -> Tuple[Message, bytes]:
         """Parse incoming message data to create a :class:`Message` instance.
 
         Any remaining message data after parsing is returned along with the instance.
+
+        Arguments:
+            msg: The incoming message bytes to parse
+            retain_text_length: Value to pass to :meth:`Display.from_dmsg`
+
+        .. versionchanged:: 0.0.8
+            The `retain_text_length` argument was added.
         """
         if len(msg) < 6:
             raise MessageParseError('Invalid header length', msg)
@@ -380,7 +421,7 @@ def parse(cls, msg: bytes) -> Tuple[Message, bytes]:
             obj.scontrol = msg
             return obj, remaining
         while len(msg):
-            disp, msg = Display.from_dmsg(obj.flags, msg)
+            disp, msg = Display.from_dmsg(obj.flags, msg, retain_text_length)
             obj.displays.append(disp)
         return obj, remaining
 

tests/conftest.py

@@ -38,6 +38,30 @@ def uhs500_msg_parsed() -> Message:
     data['displays'] = displays
     return Message(**data)
 
+
+@pytest.fixture
+def uhs500_msg_parsed_fixed_text_length() -> Message:
+    """Expected :class:`~tslumd.messages.Message` object
+    matching data from :func:`uhs500_msg_bytes`
+
+    This is the same as :func:`uhs500_msg_parsed` except that
+    the `text_length` attribute of each display is set to 12,
+    which is the fixed length used in the actual message bytes (zero-padded).
+    """
+
+    txt_length = 12
+    data = json.loads(MESSAGE_JSON.read_text())
+    data['scontrol'] = b''
+    displays = []
+    for disp in data['displays']:
+        for key in ['rh_tally', 'txt_tally', 'lh_tally']:
+            disp[key] = getattr(TallyColor, disp[key])
+        disp['text_length'] = txt_length
+        displays.append(Display(**disp))
+    data['displays'] = displays
+    return Message(**data)
+
+
 @pytest.fixture
 def udp_port0(unused_udp_port_factory):
     return unused_udp_port_factory()

tests/test_messages.py

@@ -7,18 +7,46 @@
     DmsgControlParseError, MessageParseError, MessageLengthError,
 )
 
-@pytest.fixture
-def message_with_lots_of_displays():
+
+def build_multi_display_message(num_displays: int) -> tuple[Message, list[int]]:
+    """Build a message with enough displays to require multiple packets
+    when built, returning the message object and a list of the lengths
+    of each packet.
+    """
     msgobj = Message()
-    # Message header byte length: 6
-    # Dmsg header byte length: 4
-    # Text length: 2(length bytes) + 9(chars) = 11
-    # Each Dmsg total: 4 + 11 = 15
-    # Dmsg's: 4096 * 15 = 61440
-    # 4096 Dmsg's with Message header: 4096 * 15 + 6 = 61446 bytes
-    for i in range(4096):
-        msgobj.displays.append(Display(index=i, text=f'Foo {i:05d}'))
-    return msgobj
+    msg_lengths = []
+    text_format = 'Foo {:05d}'
+    text_length = len(text_format.format(0))
+
+    # Initial message header length
+    cur_msg_length = 6
+
+    # Dmsg header + text length bytes + text bytes
+    dmsg_length = 4 + 2 + text_length
+
+    for i in range(num_displays):
+        msgobj.displays.append(Display(index=i, text=text_format.format(i)))
+        if cur_msg_length + dmsg_length > 2048:
+            msg_lengths.append(cur_msg_length)
+            cur_msg_length = 6 + dmsg_length
+        else:
+            cur_msg_length += dmsg_length
+    msg_lengths.append(cur_msg_length)
+    assert not any(l > 2048 for l in msg_lengths)
+    return msgobj, msg_lengths
+
+
+@pytest.fixture
+def message_with_multi_packet_displays() -> tuple[Message, list[int]]:
+    """Message with enough displays to require at least 6 packets
+    but not so many as to make the test take too long.
+    """
+    return build_multi_display_message(6 * 136 + 1)  # 6 full packets + 1 display
+
+
+@pytest.fixture
+def message_with_lots_of_displays() -> tuple[Message, list[int]]:
+    return build_multi_display_message(4096)
 
 
 def test_uhs_message(uhs500_msg_bytes, uhs500_msg_parsed):
@@ -60,7 +88,7 @@ def test_messages():
 
 
 def test_packet_length(faker, message_with_lots_of_displays):
-    msgobj = message_with_lots_of_displays
+    msgobj, msg_lengths = message_with_lots_of_displays
 
     # Make sure the 2048 byte limit is respected
     with pytest.raises(MessageLengthError):
@@ -73,8 +101,9 @@ def test_packet_length(faker, message_with_lots_of_displays):
 
     # Iterate over individual message packets and make sure we get all displays
     all_parsed_displays = []
-    for msg_bytes in msgobj.build_messages():
+    for i, msg_bytes in enumerate(msgobj.build_messages()):
         assert len(msg_bytes) <= 2048
+        assert len(msg_bytes) == msg_lengths[i]
         parsed, remaining = Message.parse(msg_bytes)
         assert not len(remaining)
         all_parsed_displays.extend(parsed.displays)
@@ -321,3 +350,26 @@ def test_invalid_dmsg_control(uhs500_msg_bytes, faker):
     bad_bytes = bytes(bad_bytes)
     with pytest.raises(DmsgControlParseError):
         r = Message.parse(bad_bytes)
+
+
+@pytest.mark.benchmark(group='message-parse')
+def test_bench_message_parse(uhs500_msg_bytes, uhs500_msg_parsed_fixed_text_length):
+    parsed, remaining = Message.parse(uhs500_msg_bytes, retain_text_length=True)
+    assert not len(remaining)
+    assert parsed == uhs500_msg_parsed_fixed_text_length
+
+
+@pytest.mark.benchmark(group='message-build')
+def test_bench_message_build(uhs500_msg_bytes, uhs500_msg_parsed_fixed_text_length):
+    msg_bytes = uhs500_msg_parsed_fixed_text_length.build_message()
+    assert len(msg_bytes) == len(uhs500_msg_bytes)
+    assert msg_bytes == uhs500_msg_bytes
+
+
+@pytest.mark.benchmark(group='message-build-multi')
+def test_bench_message_build_multi(message_with_multi_packet_displays):
+    msgobj, msg_lengths = message_with_multi_packet_displays
+    assert len(msg_lengths) > 1
+    for i, msg_bytes in enumerate(msgobj.build_messages()):
+        msg_len = len(msg_bytes)
+        assert msg_len == msg_lengths[i]