Generate pdu table direct. (#2752)

Author: janiversen

pymodbus/pdu/decoders.py

@@ -1,6 +1,8 @@
 """Modbus Request/Response Decoders."""
 from __future__ import annotations
 
+import copy
+
 from pymodbus.exceptions import MessageRegisterException, ModbusException
 from pymodbus.logging import Log
 
@@ -10,46 +12,44 @@
 class DecodePDU:
     """Decode pdu requests/responses (server/client)."""
 
-    _pdu_class_table: set[tuple[type[ModbusPDU], type[ModbusPDU]]] = set()
-    _pdu_sub_class_table: set[tuple[type[ModbusPDU], type[ModbusPDU]]] = set()
+    pdu_table: dict[int, tuple[type[ModbusPDU], type[ModbusPDU]]] = {}
+    pdu_sub_table: dict[int, dict[int, tuple[type[ModbusPDU], type[ModbusPDU]]]] = {}
+
 
     def __init__(self, is_server: bool) -> None:
         """Initialize function_tables."""
-        inx = 0 if is_server else 1
-        self.lookup: dict[int, type[ModbusPDU]] = {cl[inx].function_code: cl[inx] for cl in self._pdu_class_table}
-        self.sub_lookup: dict[int, dict[int, type[ModbusPDU]]] = {}
-        for f in self._pdu_sub_class_table:
-            if (function_code := f[inx].function_code) not in self.sub_lookup:
-                self.sub_lookup[function_code] = {f[inx].sub_function_code: f[inx]}
-            else:
-                self.sub_lookup[function_code][f[inx].sub_function_code] = f[inx]
+        self.pdu_inx = 0 if is_server else 1
 
     def lookupPduClass(self, data: bytes) -> type[ModbusPDU] | None:
         """Use `function_code` to determine the class of the PDU."""
-        func_code = int(data[1])
-        if func_code & 0x80:
+        if (func_code := int(data[1])) & 0x80:
             return ExceptionResponse
+        if not (pdu := self.pdu_table.get(func_code, (None, None))[self.pdu_inx]):
+            return None
+
         if func_code == 0x2B:  # mei message, sub_function_code is 1 byte
             sub_func_code = int(data[2])
-            return self.sub_lookup[func_code].get(sub_func_code, None)
-        if func_code == 0x08:  # diag message,  sub_function_code is 2 bytes
+        elif func_code == 0x08:  # diag message,  sub_function_code is 2 bytes
             sub_func_code = int.from_bytes(data[2:4], "big")
-            return self.sub_lookup[func_code].get(sub_func_code, None)
-        return self.lookup.get(func_code, None)
+        else:
+            return pdu
+        return self.pdu_sub_table[func_code].get(sub_func_code, (None, None))[self.pdu_inx]
 
     def list_function_codes(self):
         """Return list of function codes."""
-        return list(self.lookup)
+        return list(self.pdu_table)
 
     @classmethod
     def add_pdu(cls, req: type[ModbusPDU], resp: type[ModbusPDU]):
         """Register request/response."""
-        cls._pdu_class_table.add((req, resp))
+        cls.pdu_table[req.function_code] = (req, resp)
 
     @classmethod
     def add_sub_pdu(cls, req: type[ModbusPDU], resp: type[ModbusPDU]):
         """Register request/response."""
-        cls._pdu_sub_class_table.add((req, resp))
+        if req.function_code not in cls.pdu_sub_table:
+            cls.pdu_sub_table[req.function_code] = {}
+        cls.pdu_sub_table[req.function_code][req.sub_function_code] = (req, resp)
 
     def register(self, custom_class: type[ModbusPDU]) -> None:
         """Register a function and sub function class with the decoder."""
@@ -59,13 +59,9 @@ def register(self, custom_class: type[ModbusPDU]) -> None:
                 ". Class needs to be derived from "
                 "`pymodbus.pdu.ModbusPDU` "
             )
-        self.lookup[custom_class.function_code] = custom_class
-        if custom_class.sub_function_code >= 0:
-            if custom_class.function_code not in self.sub_lookup:
-                self.sub_lookup[custom_class.function_code] = {}
-            self.sub_lookup[custom_class.function_code][
-                custom_class.sub_function_code
-            ] = custom_class
+        if "pdu_table" not in self.__dict__:
+            self.pdu_table = copy.deepcopy(DecodePDU.pdu_table)
+        self.pdu_table[custom_class.function_code] = (custom_class, custom_class)
 
     def decode(self, frame: bytes) -> ModbusPDU | None:
         """Decode a frame."""
@@ -74,14 +70,14 @@ def decode(self, frame: bytes) -> ModbusPDU | None:
                 pdu_exp = ExceptionResponse(function_code & 0x7F)
                 pdu_exp.decode(frame[1:])
                 return pdu_exp
-            if not (pdu_class := self.lookup.get(function_code, None)):
+            if not (pdu_class := self.pdu_table.get(function_code, (None, None))[self.pdu_inx]):
                 Log.debug("decode PDU failed for function code {}", function_code)
                 raise ModbusException(f"Unknown response {function_code}")
             pdu = pdu_class()
             pdu.decode(frame[1:])
             if pdu.sub_function_code >= 0:
-                lookup = self.sub_lookup.get(pdu.function_code, {})
-                if sub_class := lookup.get(pdu.sub_function_code, None):
+                lookup = self.pdu_sub_table.get(pdu.function_code, {})
+                if sub_class := lookup.get(pdu.sub_function_code, (None,None))[self.pdu_inx]:
                     pdu = sub_class()
                     pdu.decode(frame[1:])
             Log.debug("decoded PDU function_code({} sub {}) -> {} ", pdu.function_code, pdu.sub_function_code, str(pdu))