Clean up construct configuration for nibe gw (#179)

This reduces the nesting depth and make the logged output from a parsed block more readable.
yozik04 · Oct 13, 2024 · 9431e9a · 9431e9a
2 parents 4d373f5 + 2195a07
commit 9431e9a
Show file tree

Hide file tree

Showing 3 changed files with 162 additions and 132 deletions.
diff --git a/nibe/connection/nibegw.py b/nibe/connection/nibegw.py
@@ -15,7 +15,6 @@
 
 from construct import (
     Adapter,
-    Array,
     BitStruct,
     Bytes,
     Checksum,
@@ -24,10 +23,11 @@
     Container,
     Enum,
     EnumIntegerString,
-    FixedSized,
     Flag,
     FlagsEnum,
+    FocusedSeq,
     GreedyBytes,
+    GreedyRange,
     GreedyString,
     IfThenElse,
     Int8sb,
@@ -36,6 +36,7 @@
     Int16ub,
     Int16ul,
     NullTerminated,
+    Peek,
     Pointer,
     Prefixed,
     RawCopy,
@@ -584,6 +585,35 @@ def _encode(self, obj, context, path):
         return obj / self._scale - self._offset
 
 
+StartCode = Enum(
+    Int8ub,
+    RESPONSE=0x5C,
+    REQUEST=0xC0,
+    ACK=0x06,
+    NAK=0x15,
+)
+
+Command = Enum(
+    Int8ub,
+    RMU_WRITE_REQ=0x60,
+    RMU_DATA_MSG=0x62,
+    RMU_DATA_REQ=0x63,
+    MODBUS_DATA_MSG=0x68,
+    MODBUS_READ_REQ=0x69,
+    MODBUS_READ_RESP=0x6A,
+    MODBUS_WRITE_REQ=0x6B,
+    MODBUS_WRITE_RESP=0x6C,
+    MODBUS_ADDRESS_MSG=0x6E,
+    PRODUCT_INFO_MSG=0x6D,
+    ACCESSORY_VERSION_REQ=0xEE,
+    ECS_DATA_REQ=0x90,
+    ECS_DATA_MSG_1=0x55,
+    ECS_DATA_MSG_2=0xA0,
+    STRING_MSG=0xB1,
+    HEATPUMP_REQ=0xF7,
+)
+
+
 StringData = Struct(
     "unknown" / Int8ub,
     "id" / Int16ul,
@@ -655,45 +685,21 @@ def _encode(self, obj, context, path):
     "unknown5" / GreedyBytes,
 )
 
-Data = Dedupe5C(
-    Switch(
-        this.cmd,
-        {
-            "MODBUS_READ_RESP": Struct("coil_address" / Int16ul, "value" / Bytes(4)),
-            "MODBUS_DATA_MSG": Array(
-                lambda this: this.length // 4,
-                Struct("coil_address" / Int16ul, "value" / Bytes(2)),
-            ),
-            "MODBUS_WRITE_RESP": Struct("result" / Flag),
-            "MODBUS_ADDRESS_MSG": Struct("address" / Int8ub),
-            "PRODUCT_INFO_MSG": ProductInfoData,
-            "RMU_DATA_MSG": RmuData,
-            "STRING_MSG": StringData,
-        },
-        default=Bytes(this.length),
-    )
-)
-
-
-Command = Enum(
-    Int8ub,
-    RMU_WRITE_REQ=0x60,
-    RMU_DATA_MSG=0x62,
-    RMU_DATA_REQ=0x63,
-    MODBUS_DATA_MSG=0x68,
-    MODBUS_READ_REQ=0x69,
-    MODBUS_READ_RESP=0x6A,
-    MODBUS_WRITE_REQ=0x6B,
-    MODBUS_WRITE_RESP=0x6C,
-    MODBUS_ADDRESS_MSG=0x6E,
-    PRODUCT_INFO_MSG=0x6D,
-    ACCESSORY_VERSION_REQ=0xEE,
-    ECS_DATA_REQ=0x90,
-    ECS_DATA_MSG_1=0x55,
-    ECS_DATA_MSG_2=0xA0,
-    STRING_MSG=0xB1,
-    HEATPUMP_REQ=0xF7,
-)
+ModbusDataValue = Struct("coil_address" / Int16ul, "value" / Bytes(2))
+ModbusData = GreedyRange(ModbusDataValue)
+ModbusReadResp = Struct("coil_address" / Int16ul, "value" / Bytes(4))
+ModbusWriteResp = Struct("result" / Flag)
+ModbusAddressMsg = Struct("address" / Int8ub)
+
+ResponseTypes = {
+    "MODBUS_READ_RESP": ModbusReadResp,
+    "MODBUS_DATA_MSG": ModbusData,
+    "MODBUS_WRITE_RESP": ModbusWriteResp,
+    "MODBUS_ADDRESS_MSG": ModbusAddressMsg,
+    "PRODUCT_INFO_MSG": ProductInfoData,
+    "RMU_DATA_MSG": RmuData,
+    "STRING_MSG": StringData,
+}
 
 Address = Enum(
     Int16ub,
@@ -736,80 +742,128 @@ def _encode(self, obj, context, path):
 
 
 # fmt: off
-Response = Struct(
-    "start_byte" / Const(0x5C, Int8ub),
-    "fields" / RawCopy(
-        Struct(
-            "address" / Address,
-            "cmd" / Command,
-            "length" / Int8ub,
-            "data" / FixedSized(this.length, Data),
+
+ResponseData = Struct(
+    "address" / Address,
+    "cmd" / Command,
+    "data" / Prefixed(Int8ub,
+        Dedupe5C(
+            Switch(
+                this.cmd, ResponseTypes,
+                default=GreedyBytes,
+            )
         )
     ),
+)
+
+Response = Struct(
+    "start_byte" / Const("RESPONSE", StartCode),
+    "fields" / RawCopy(ResponseData),
     "checksum" / Checksum(Int8ub, xor8, this.fields.data),
 )
 
+AccessoryVersionReq = UnionConstruct(None,
+    # Modbus and RMU seem to disagree on how to interpret this
+    # data, at least from how it looks in the service info screen
+    # on the pump.
+    "modbus" / Struct(
+        "version" / Int16ul,
+        "unknown" / Int8ub,
+    ),
+    "rmu" / Struct(
+        "unknown" / Int8ub,
+        "version" / Int16ul,
+    ),
+)
+
+RmuWriteReqTypes = {
+    "TEMPORARY_LUX": Int8ub,
+    "TEMPERATURE": FixedPoint(Int16ul, 0.1, -7.0, size="s16"),
+    "FUNCTIONS": FlagsEnum(
+        Int8ub,
+        allow_additive_heating=0x01,
+        allow_heating=0x02,
+        allow_cooling=0x04,
+    ),
+    "OPERATIONAL_MODE": Int8ub,
+    "SETPOINT_S1": FixedPoint(Int16sl, 0.1, 0.0, size="s16"),
+    "SETPOINT_S2": FixedPoint(Int16sl, 0.1, 0.0, size="s16"),
+    "SETPOINT_S3": FixedPoint(Int16sl, 0.1, 0.0, size="s16"),
+    "SETPOINT_S4": FixedPoint(Int16sl, 0.1, 0.0, size="s16"),
+}
 
-RequestData = Switch(
-    this.cmd,
-    {
-        "ACCESSORY_VERSION_REQ": UnionConstruct(None,
-            # Modbus and RMU seem to disagree on how to interpret this
-            # data, at least from how it looks in the service info screen
-            # on the pump.
-            "modbus" / Struct(
-                "version" / Int16ul,
-                "unknown" / Int8ub,
-            ),
-            "rmu" / Struct(
-                "unknown" / Int8ub,
-                "version" / Int16ul,
-            ),
-        ),
-        "RMU_WRITE_REQ": Struct(
-            "index" / RmuWriteIndex,
-            "value" / Switch(
-                lambda this: this.index,
-                {
-                    "TEMPORARY_LUX": Int8ub,
-                    "TEMPERATURE": FixedPoint(Int16ul, 0.1, -7.0, size="s16"),
-                    "FUNCTIONS": FlagsEnum(
-                        Int8ub,
-                        allow_additive_heating=0x01,
-                        allow_heating=0x02,
-                        allow_cooling=0x04,
-                    ),
-                    "OPERATIONAL_MODE": Int8ub,
-                    "SETPOINT_S1": FixedPoint(Int16sl, 0.1, 0.0, size="s16"),
-                    "SETPOINT_S2": FixedPoint(Int16sl, 0.1, 0.0, size="s16"),
-                    "SETPOINT_S3": FixedPoint(Int16sl, 0.1, 0.0, size="s16"),
-                    "SETPOINT_S4": FixedPoint(Int16sl, 0.1, 0.0, size="s16"),
-                },
-                default=Select(
-                    Int16ul,
-                    Int8ub,
-                )
-            )
-        ),
-        "MODBUS_READ_REQ": Struct(
-            "coil_address" / Int16ul,
-        ),
-        "MODBUS_WRITE_REQ": Struct(
-            "coil_address" / Int16ul,
-            "value" / Bytes(4),
+RmuWriteReq = Struct(
+    "index" / RmuWriteIndex,
+    "value" / Switch(
+        this.index,
+        RmuWriteReqTypes,
+        default=Select(
+            Int16ul,
+            Int8ub,
+        )
+    )
+)
+
+ModbusReadReq = Struct(
+    "coil_address" / Int16ul,
+)
+
+ModbusWriteReq = Struct(
+    "coil_address" / Int16ul,
+    "value" / Bytes(4),
+)
+
+RequestTypes = {
+    "ACCESSORY_VERSION_REQ": AccessoryVersionReq,
+    "RMU_WRITE_REQ": RmuWriteReq,
+    "MODBUS_READ_REQ": ModbusReadReq,
+    "MODBUS_WRITE_REQ": ModbusWriteReq
+}
+
+RequestData = Struct(
+    Const("REQUEST", StartCode),
+    "cmd" / Command,
+    "data" / Prefixed(Int8ub,
+        Switch(
+            this.cmd,
+            RequestTypes,
+            default=GreedyBytes,
         )
-    },
-    default=Bytes(this.length),
+    )
 )
 
 Request = Struct(
-    "fields" / RawCopy(
-        Struct(
-            "start_byte" / Const(0xC0, Int8ub),
-            "cmd" / Command,
-            "data" / Prefixed(Int8ub, RequestData)
-        )
-    ),
+    "start_byte" / Peek(StartCode),
+    "fields" / RawCopy(RequestData),
     "checksum" / Checksum(Int8ub, xor8, this.fields.data),
 )
+
+AckData = Struct(Const("ACK", StartCode))
+Ack = Struct(
+    "start_byte" / Peek(StartCode),
+    "fields" / RawCopy(AckData)
+)
+
+NakData = Struct(Const("NAK", StartCode))
+Nak = Struct(
+    "start_byte" / Peek(StartCode),
+    "fields" / RawCopy(NakData)
+)
+
+BlockTypes = {
+    "RESPONSE": Response,
+    "REQUEST": Request,
+    "ACK": Ack,
+    "NAK": Nak,
+}
+
+Block = FocusedSeq(
+    "data",
+    "start" / Peek(StartCode),
+    "data" / Switch(
+        this.start,
+        BlockTypes
+    ),
+)
+
 # fmt: on
diff --git a/nibe/console_scripts/cli.py b/nibe/console_scripts/cli.py
@@ -9,38 +9,15 @@
 from typing import IO
 
 import asyncclick as click
-from construct import (
-    Const,
-    ConstructError,
-    FocusedSeq,
-    GreedyRange,
-    Int8ul,
-    Peek,
-    RawCopy,
-    Struct,
-    Switch,
-    this,
-)
+from construct import ConstructError
 
 from ..coil import CoilData
 from ..connection import Connection
 from ..connection.modbus import Modbus
-from ..connection.nibegw import NibeGW, Request, Response
+from ..connection.nibegw import Block, NibeGW
 from ..exceptions import NibeException
 from ..heatpump import HeatPump, Model
 
-Ack = Struct("fields" / RawCopy(Struct("Ack" / Const(0x06, Int8ul))))
-
-Nak = Struct("fields" / RawCopy(Struct("Nak" / Const(0x15, Int8ul))))
-
-Block = FocusedSeq(
-    "data",
-    "start" / Peek(Int8ul),
-    "data" / Switch(this.start, {0x5C: Response, 0xC0: Request, 0x06: Ack, 0x15: Nak}),
-)
-
-Stream = GreedyRange(Block)
-
 
 @click.group()
 @click.option("-v", "--verbose", count=True)
@@ -246,7 +223,7 @@ def parse_file(file: IO, type: str):
 
     with io.BytesIO(bytes(data)) as stream:
         for packet in parse_stream(stream):
-            click.echo(packet.fields.value)
+            click.echo(packet)
 
         remaining = stream.read()
         if remaining:

diff --git a/tests/connection/test_nibegw_message_parsing.py b/tests/connection/test_nibegw_message_parsing.py
@@ -20,7 +20,6 @@ def test_parse_token_response_16bit_address(self):
         data = self._parse_hexlified_raw_message("5c41c9f7007f06")
         assert data.address == "HEATPUMP_1"
         assert data.cmd == "HEATPUMP_REQ"
-        assert data.length == 0
         assert data.data == b""
 
     def test_parse_escaped_read_response(self):