feat: migrate filter diags and leadmessage to pydantic-xml

cryptk · cryptk · commit 7a37b1da7e21 · 2025-11-05T22:44:49.000-06:00
diff --git a/pyomnilogic_local/api/protocol.py b/pyomnilogic_local/api/protocol.py
@@ -363,7 +363,7 @@ async def _receive_file(self) -> str:
         # If the response is too large, the controller will send a LeadMessage indicating how many follow-up messages will be sent
         if message.type is MessageType.MSP_LEADMESSAGE:
             try:
-                leadmsg = LeadMessage.model_validate(ET.fromstring(message.payload[:-1]))
+                leadmsg = LeadMessage.from_xml(message.payload[:-1])
             except Exception as exc:
                 raise OmniFragmentationException(f"Failed to parse LeadMessage: {exc}") from exc
 
diff --git a/pyomnilogic_local/cli/pcap_utils.py b/pyomnilogic_local/cli/pcap_utils.py
@@ -7,7 +7,6 @@
 
 from __future__ import annotations
 
-import xml.etree.ElementTree as ET
 import zlib
 from collections import defaultdict
 from typing import Any
@@ -147,7 +146,7 @@ def process_pcap_messages(packets: Any) -> list[tuple[str, str, OmniLogicMessage
 
                 # Check if we have all the blocks
                 lead_msg = message_sequences[matching_seq][0]
-                lead_data = LeadMessage.model_validate(ET.fromstring(lead_msg.payload[:-1]))
+                lead_data = LeadMessage.from_xml(lead_msg.payload[:-1])
 
                 # We have LeadMessage + all BlockMessages
                 if len(message_sequences[matching_seq]) == lead_data.msg_block_count + 1:
diff --git a/pyomnilogic_local/models/filter_diagnostics.py b/pyomnilogic_local/models/filter_diagnostics.py
@@ -1,42 +1,87 @@
 from __future__ import annotations
 
-from pydantic import BaseModel, ConfigDict, Field
-from xmltodict import parse as xml_parse
+from pydantic import ConfigDict
+from pydantic_xml import BaseXmlModel, attr, element, wrapped
 
+from .const import XML_NS
 
-class FilterDiagnosticsParameter(BaseModel):
-    model_config = ConfigDict(from_attributes=True)
+# Example Filter Diagnostics XML:
+#
+# <?xml version="1.0" encoding="UTF-8" ?>
+# <Response xmlns="http://nextgen.hayward.com/api">
+#     <Name>GetUIFilterDiagnosticInfoRsp</Name>
+#     <Parameters>
+#         <Parameter name="PoolID" dataType="int">7</Parameter>
+#         <Parameter name="EquipmentID" dataType="int">8</Parameter>
+#         <Parameter name="PowerLSB" dataType="byte">133</Parameter>
+#         <Parameter name="PowerMSB" dataType="byte">4</Parameter>
+#         <Parameter name="ErrorStatus" dataType="byte">0</Parameter>
+#         <Parameter name="DisplayFWRevisionB1" dataType="byte">49</Parameter>
+#         <Parameter name="DisplayFWRevisionB2" dataType="byte">48</Parameter>
+#         <Parameter name="DisplayFWRevisionB3" dataType="byte">49</Parameter>
+#         <Parameter name="DisplayFWRevisionB4" dataType="byte">53</Parameter>
+#         <Parameter name="DisplayFWRevisionB5" dataType="byte">32</Parameter>
+#         <Parameter name="DisplayFWRevisionB6" dataType="byte">0</Parameter>
+#         <Parameter name="DriveFWRevisionB1" dataType="byte">48</Parameter>
+#         <Parameter name="DriveFWRevisionB2" dataType="byte">48</Parameter>
+#         <Parameter name="DriveFWRevisionB3" dataType="byte">55</Parameter>
+#         <Parameter name="DriveFWRevisionB4" dataType="byte">48</Parameter>
+#         <Parameter name="DriveFWRevisionB5" dataType="byte">32</Parameter>
+#         <Parameter name="DriveFWRevisionB6" dataType="byte">0</Parameter>
+#     </Parameters>
+# </Response>
 
-    name: str = Field(alias="@name")
-    dataType: str = Field(alias="@dataType")
-    value: int = Field(alias="#text")
 
+class FilterDiagnosticsParameter(BaseXmlModel, tag="Parameter", ns="api", nsmap=XML_NS):
+    """Individual diagnostic parameter with name, type, and value."""
 
-class FilterDiagnosticsParameters(BaseModel):
     model_config = ConfigDict(from_attributes=True)
 
-    parameter: list[FilterDiagnosticsParameter] = Field(alias="Parameter")
+    name: str = attr()
+    data_type: str = attr(name="dataType")
+    value: int
+
 
+class FilterDiagnostics(BaseXmlModel, tag="Response", ns="api", nsmap=XML_NS):
+    """Filter diagnostics response containing diagnostic parameters.
+
+    The XML structure has a Parameters wrapper element containing Parameter children:
+    <Response>
+        <Name>FilterDiagnostics</Name>
+        <Parameters>
+            <Parameter name="..." dataType="...">value</Parameter>
+            ...
+        </Parameters>
+    </Response>
+    """
 
-class FilterDiagnostics(BaseModel):
     model_config = ConfigDict(from_attributes=True)
 
-    name: str = Field(alias="Name")
-    # parameters: FilterDiagnosticsParameters = Field(alias="Parameters")
-    parameters: list[FilterDiagnosticsParameter] = Field(alias="Parameters")
+    name: str = element(tag="Name")
+    parameters: list[FilterDiagnosticsParameter] = wrapped("Parameters", element(tag="Parameter", default_factory=list))
 
     def get_param_by_name(self, name: str) -> int:
+        """Get parameter value by name.
+
+        Args:
+            name: Name of the parameter to retrieve
+
+        Returns:
+            The integer value of the parameter
+
+        Raises:
+            IndexError: If parameter name not found
+        """
         return [param.value for param in self.parameters if param.name == name][0]
 
     @staticmethod
     def load_xml(xml: str) -> FilterDiagnostics:
-        data = xml_parse(
-            xml,
-            # Some things will be lists or not depending on if a pool has more than one of that piece of equipment.  Here we are coercing
-            # everything that *could* be a list into a list to make the parsing more consistent.
-            force_list=("Parameter"),
-        )
-        # The XML nests the Parameter entries under a Parameters entry, this is annoying to work with.  Here we are adjusting the data to
-        # remove that extra level in the data
-        data["Response"]["Parameters"] = data["Response"]["Parameters"]["Parameter"]
-        return FilterDiagnostics.model_validate(data["Response"])
+        """Load filter diagnostics from XML string.
+
+        Args:
+            xml: XML string containing filter diagnostics data
+
+        Returns:
+            Parsed FilterDiagnostics instance
+        """
+        return FilterDiagnostics.from_xml(xml)
diff --git a/pyomnilogic_local/models/leadmessage.py b/pyomnilogic_local/models/leadmessage.py
@@ -1,30 +1,63 @@
 from __future__ import annotations
 
-from typing import Any
-from xml.etree.ElementTree import Element
-
-from pydantic import BaseModel, ConfigDict, Field, model_validator
+from pydantic import ConfigDict, computed_field
+from pydantic_xml import BaseXmlModel, attr, element, wrapped
 
 from .const import XML_NS
 
+# Example Lead Message XML:
+#
+# <?xml.version="1.0" encoding="UTF-8"?>
+# <Response xmlns="http://nextgen.hayward.com/api">
+#     <Name>LeadMessage</Name>
+#     <Parameters>
+#         <Parameter name="SourceOpId" dataType="int">1003</Parameter>
+#         <Parameter name="MsgSize" dataType="int">3709</Parameter>
+#         <Parameter name="MsgBlockCount" dataType="int">4</Parameter>
+#         <Parameter name="Type" dataType="int">0</Parameter>
+#     </Parameters>
+# </Response>
+
+
+class LeadMessageParameter(BaseXmlModel, tag="Parameter", ns="api", nsmap=XML_NS):
+    """Individual parameter in lead message."""
+
+    name: str = attr()
+    value: int
+
+
+class LeadMessage(BaseXmlModel, tag="Response", ns="api", nsmap=XML_NS):
+    """Lead message containing protocol parameters.
+
+    Lead messages are sent at the start of communication to establish
+    protocol parameters like message size and block count.
+    """
 
-class LeadMessage(BaseModel):
     model_config = ConfigDict(from_attributes=True)
 
-    source_op_id: int = Field(alias="SourceOpId")
-    msg_size: int = Field(alias="MsgSize")
-    msg_block_count: int = Field(alias="MsgBlockCount")
-    type: int = Field(alias="Type")
-
-    @model_validator(mode="before")
-    @classmethod
-    def parse_xml_element(cls, data: Any) -> dict[str, Any]:
-        """Parse XML Element into dict format for Pydantic validation."""
-        if isinstance(data, Element):
-            # Parse the Parameter elements from the XML
-            result = {}
-            for param in data.findall(".//api:Parameter", XML_NS):
-                if name := param.get("name"):
-                    result[name] = int(param.text) if param.text else 0
-            return result
-        return data
+    name: str = element(tag="Name")
+    parameters: list[LeadMessageParameter] = wrapped("Parameters", element(tag="Parameter", default_factory=list))
+
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def source_op_id(self) -> int:
+        """Extract SourceOpId from parameters."""
+        return next((p.value for p in self.parameters if p.name == "SourceOpId"), 0)
+
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def msg_size(self) -> int:
+        """Extract MsgSize from parameters."""
+        return next((p.value for p in self.parameters if p.name == "MsgSize"), 0)
+
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def msg_block_count(self) -> int:
+        """Extract MsgBlockCount from parameters."""
+        return next((p.value for p in self.parameters if p.name == "MsgBlockCount"), 0)
+
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def type(self) -> int:
+        """Extract Type from parameters."""
+        return next((p.value for p in self.parameters if p.name == "Type"), 0)
diff --git a/tests/test_fixtures.py b/tests/test_fixtures.py
@@ -109,7 +109,7 @@ def test_mspconfig(self, fixture_data: dict[str, str], subtests: SubTests) -> No
             pool = msp.backyard.bow[0]
             assert pool.system_id == 3
             assert pool.name == "Pool"
-            assert pool.omni_type == OmniType.BOW_MSP
+            assert pool.omni_type == OmniType.BOW
 
         with subtests.test(msg="filter configuration"):
             filters = get_equipment_by_type(msp, OmniType.FILTER)
@@ -155,7 +155,7 @@ def test_telemetry(self, fixture_data: dict[str, str], subtests: SubTests) -> No
             assert len(telem.virtual_heater) == 1
             vh = telem.virtual_heater[0]
             assert vh.system_id == 15
-            assert vh.current_set_point == 82
+            assert vh.current_set_point == 65
 
         with subtests.test(msg="heater equipment telemetry"):
             assert telem.heater is not None
@@ -202,7 +202,7 @@ def test_mspconfig(self, fixture_data: dict[str, str], subtests: SubTests) -> No
             pool = msp.backyard.bow[0]
             assert pool.system_id == 10
             assert pool.name == "Pool"
-            assert pool.omni_type == OmniType.BOW_MSP
+            assert pool.omni_type == OmniType.BOW
 
         with subtests.test(msg="filter configuration"):
             filters = get_equipment_by_type(msp, OmniType.FILTER)
@@ -225,7 +225,7 @@ def test_mspconfig(self, fixture_data: dict[str, str], subtests: SubTests) -> No
         with subtests.test(msg="backyard sensors"):
             assert msp.backyard.sensor is not None
             assert len(msp.backyard.sensor) == 1
-            assert msp.backyard.sensor[0].system_id == 1
+            assert msp.backyard.sensor[0].system_id == 16
 
     def test_telemetry(self, fixture_data: dict[str, str], subtests: SubTests) -> None:
         """Test Telemetry parsing for issue-163."""
@@ -299,7 +299,7 @@ def test_mspconfig(self, fixture_data: dict[str, str], subtests: SubTests) -> No
             pool = msp.backyard.bow[0]
             assert pool.system_id == 1
             assert pool.name == "Pool"
-            assert pool.omni_type == OmniType.BOW_MSP
+            assert pool.omni_type == OmniType.BOW
 
         with subtests.test(msg="spa configuration"):
             assert msp.backyard.bow is not None
@@ -378,7 +378,7 @@ def test_telemetry(self, fixture_data: dict[str, str], subtests: SubTests) -> No
 
         with subtests.test(msg="relay telemetry"):
             assert telem.relay is not None
-            assert len(telem.relay) == 2
+            assert len(telem.relay) == 3
             # Check yard lights relay is on
             yard_relay = [r for r in telem.relay if r.system_id == 27][0]
             assert yard_relay.state.value == 1  # ON
