Convert if/elif chains to dicts in modbus (#51962)

homeassistant/components/modbus/__init__.py

@@ -69,7 +69,9 @@ from .const import (
     CONF_RETRIES,
     CONF_RETRY_ON_EMPTY,
     CONF_REVERSE_ORDER,
+    CONF_RTUOVERTCP,
     CONF_SCALE,
+    CONF_SERIAL,
     CONF_STATE_CLOSED,
     CONF_STATE_CLOSING,
     CONF_STATE_OFF,
@@ -86,6 +88,8 @@ from .const import (
     CONF_SWAP_WORD,
     CONF_SWAP_WORD_BYTE,
     CONF_TARGET_TEMP,
+    CONF_TCP,
+    CONF_UDP,
     CONF_VERIFY,
     CONF_WRITE_TYPE,
     DATA_TYPE_CUSTOM,
@@ -292,7 +296,7 @@ MODBUS_SCHEMA = vol.Schema(
 
 SERIAL_SCHEMA = MODBUS_SCHEMA.extend(
     {
-        vol.Required(CONF_TYPE): "serial",
+        vol.Required(CONF_TYPE): CONF_SERIAL,
         vol.Required(CONF_BAUDRATE): cv.positive_int,
         vol.Required(CONF_BYTESIZE): vol.Any(5, 6, 7, 8),
         vol.Required(CONF_METHOD): vol.Any("rtu", "ascii"),
@@ -306,7 +310,7 @@ ETHERNET_SCHEMA = MODBUS_SCHEMA.extend(
     {
         vol.Required(CONF_HOST): cv.string,
         vol.Required(CONF_PORT): cv.port,
-        vol.Required(CONF_TYPE): vol.Any("tcp", "udp", "rtuovertcp"),
+        vol.Required(CONF_TYPE): vol.Any(CONF_TCP, CONF_UDP, CONF_RTUOVERTCP),
     }
 )
 

homeassistant/components/modbus/const.py

@@ -39,7 +39,9 @@ CONF_RETRIES = "retries"
 CONF_RETRY_ON_EMPTY = "retry_on_empty"
 CONF_REVERSE_ORDER = "reverse_order"
 CONF_PRECISION = "precision"
+CONF_RTUOVERTCP = "rtuovertcp"
 CONF_SCALE = "scale"
+CONF_SERIAL = "serial"
 CONF_STATE_CLOSED = "state_closed"
 CONF_STATE_CLOSING = "state_closing"
 CONF_STATE_OFF = "state_off"
@@ -56,6 +58,8 @@ CONF_SWAP_NONE = "none"
 CONF_SWAP_WORD = "word"
 CONF_SWAP_WORD_BYTE = "word_byte"
 CONF_TARGET_TEMP = "target_temp_register"
+CONF_TCP = "tcp"
+CONF_UDP = "udp"
 CONF_VERIFY = "verify"
 CONF_VERIFY_REGISTER = "verify_register"
 CONF_VERIFY_STATE = "verify_state"

homeassistant/components/modbus/modbus.py

@@ -5,7 +5,6 @@ import logging
 from pymodbus.client.sync import ModbusSerialClient, ModbusTcpClient, ModbusUdpClient
 from pymodbus.constants import Defaults
 from pymodbus.exceptions import ModbusException
-from pymodbus.transaction import ModbusRtuFramer
 
 from homeassistant.const import (
     CONF_DELAY,
@@ -41,7 +40,11 @@ from .const import (
     CONF_PARITY,
     CONF_RETRIES,
     CONF_RETRY_ON_EMPTY,
+    CONF_RTUOVERTCP,
+    CONF_SERIAL,
     CONF_STOPBITS,
+    CONF_TCP,
+    CONF_UDP,
     DEFAULT_HUB,
     MODBUS_DOMAIN as DOMAIN,
     PLATFORMS,
@@ -51,9 +54,53 @@ from .const import (
 
 ENTRY_FUNC = "func"
 ENTRY_ATTR = "attr"
+ENTRY_NAME = "name"
 
 _LOGGER = logging.getLogger(__name__)
 
+PYMODBUS_CALL = {
+    CALL_TYPE_COIL: {
+        ENTRY_ATTR: "bits",
+        ENTRY_NAME: "read_coils",
+        ENTRY_FUNC: None,
+    },
+    CALL_TYPE_DISCRETE: {
+        ENTRY_ATTR: "bits",
+        ENTRY_NAME: "read_discrete_inputs",
+        ENTRY_FUNC: None,
+    },
+    CALL_TYPE_REGISTER_HOLDING: {
+        ENTRY_ATTR: "registers",
+        ENTRY_NAME: "read_holding_registers",
+        ENTRY_FUNC: None,
+    },
+    CALL_TYPE_REGISTER_INPUT: {
+        ENTRY_ATTR: "registers",
+        ENTRY_NAME: "read_input_registers",
+        ENTRY_FUNC: None,
+    },
+    CALL_TYPE_WRITE_COIL: {
+        ENTRY_ATTR: "value",
+        ENTRY_NAME: "write_coil",
+        ENTRY_FUNC: None,
+    },
+    CALL_TYPE_WRITE_COILS: {
+        ENTRY_ATTR: "count",
+        ENTRY_NAME: "write_coils",
+        ENTRY_FUNC: None,
+    },
+    CALL_TYPE_WRITE_REGISTER: {
+        ENTRY_ATTR: "value",
+        ENTRY_NAME: "write_register",
+        ENTRY_FUNC: None,
+    },
+    CALL_TYPE_WRITE_REGISTERS: {
+        ENTRY_ATTR: "count",
+        ENTRY_NAME: "write_registers",
+        ENTRY_FUNC: None,
+    },
+}
+
 
 async def async_modbus_setup(
     hass, config, service_write_register_schema, service_write_coil_schema
@@ -147,58 +194,39 @@ class ModbusHub:
         self.hass = hass
         self._config_name = client_config[CONF_NAME]
         self._config_type = client_config[CONF_TYPE]
-        self._config_port = client_config[CONF_PORT]
-        self._config_timeout = client_config[CONF_TIMEOUT]
         self._config_delay = client_config[CONF_DELAY]
-        self._config_reset_socket = client_config[CONF_CLOSE_COMM_ON_ERROR]
+        self._pb_call = PYMODBUS_CALL.copy()
-        self._config_retries = client_config[CONF_RETRIES]
+        self._pb_class = {
-        self._config_retry_on_empty = client_config[CONF_RETRY_ON_EMPTY]
+            CONF_SERIAL: ModbusSerialClient,
-        Defaults.Timeout = client_config[CONF_TIMEOUT]
+            CONF_TCP: ModbusTcpClient,
-        if self._config_type == "serial":
+            CONF_UDP: ModbusUdpClient,
+            CONF_RTUOVERTCP: ModbusTcpClient,
+        }
+        self._pb_params = {
+            "port": client_config[CONF_PORT],
+            "timeout": client_config[CONF_TIMEOUT],
+            "reset_socket": client_config[CONF_CLOSE_COMM_ON_ERROR],
+            "retries": client_config[CONF_RETRIES],
+            "retry_on_empty": client_config[CONF_RETRY_ON_EMPTY],
+        }
+        if self._config_type == CONF_SERIAL:
             # serial configuration
-            self._config_method = client_config[CONF_METHOD]
+            self._pb_params.update(
-            self._config_baudrate = client_config[CONF_BAUDRATE]
+                {
-            self._config_stopbits = client_config[CONF_STOPBITS]
+                    "method": client_config[CONF_METHOD],
-            self._config_bytesize = client_config[CONF_BYTESIZE]
+                    "baudrate": client_config[CONF_BAUDRATE],
-            self._config_parity = client_config[CONF_PARITY]
+                    "stopbits": client_config[CONF_STOPBITS],
+                    "bytesize": client_config[CONF_BYTESIZE],
+                    "parity": client_config[CONF_PARITY],
+                }
+            )
         else:
             # network configuration
-            self._config_host = client_config[CONF_HOST]
+            self._pb_params["host"] = client_config[CONF_HOST]
+            if self._config_type == CONF_RTUOVERTCP:
+                self._pb_params["host"] = "ModbusRtuFramer"
 
-        self._call_type = {
+        Defaults.Timeout = client_config[CONF_TIMEOUT]
-            CALL_TYPE_COIL: {
-                ENTRY_ATTR: "bits",
-                ENTRY_FUNC: None,
-            },
-            CALL_TYPE_DISCRETE: {
-                ENTRY_ATTR: "bits",
-                ENTRY_FUNC: None,
-            },
-            CALL_TYPE_REGISTER_HOLDING: {
-                ENTRY_ATTR: "registers",
-                ENTRY_FUNC: None,
-            },
-            CALL_TYPE_REGISTER_INPUT: {
-                ENTRY_ATTR: "registers",
-                ENTRY_FUNC: None,
-            },
-            CALL_TYPE_WRITE_COIL: {
-                ENTRY_ATTR: "value",
-                ENTRY_FUNC: None,
-            },
-            CALL_TYPE_WRITE_COILS: {
-                ENTRY_ATTR: "count",
-                ENTRY_FUNC: None,
-            },
-            CALL_TYPE_WRITE_REGISTER: {
-                ENTRY_ATTR: "value",
-                ENTRY_FUNC: None,
-            },
-            CALL_TYPE_WRITE_REGISTERS: {
-                ENTRY_ATTR: "count",
-                ENTRY_FUNC: None,
-            },
-        }
 
     def _log_error(self, text: str, error_state=True):
         log_text = f"Pymodbus: {text}"
@@ -211,75 +239,19 @@ class ModbusHub:
     async def async_setup(self):
         """Set up pymodbus client."""
         try:
-            if self._config_type == "serial":
+            self._client = self._pb_class[self._config_type](**self._pb_params)
-                self._client = ModbusSerialClient(
-                    method=self._config_method,
-                    port=self._config_port,
-                    baudrate=self._config_baudrate,
-                    stopbits=self._config_stopbits,
-                    bytesize=self._config_bytesize,
-                    parity=self._config_parity,
-                    timeout=self._config_timeout,
-                    retries=self._config_retries,
-                    retry_on_empty=self._config_retry_on_empty,
-                    reset_socket=self._config_reset_socket,
-                )
-            elif self._config_type == "rtuovertcp":
-                self._client = ModbusTcpClient(
-                    host=self._config_host,
-                    port=self._config_port,
-                    framer=ModbusRtuFramer,
-                    timeout=self._config_timeout,
-                    retries=self._config_retries,
-                    retry_on_empty=self._config_retry_on_empty,
-                    reset_socket=self._config_reset_socket,
-                )
-            elif self._config_type == "tcp":
-                self._client = ModbusTcpClient(
-                    host=self._config_host,
-                    port=self._config_port,
-                    timeout=self._config_timeout,
-                    retries=self._config_retries,
-                    retry_on_empty=self._config_retry_on_empty,
-                    reset_socket=self._config_reset_socket,
-                )
-            elif self._config_type == "udp":
-                self._client = ModbusUdpClient(
-                    host=self._config_host,
-                    port=self._config_port,
-                    timeout=self._config_timeout,
-                    retries=self._config_retries,
-                    retry_on_empty=self._config_retry_on_empty,
-                    reset_socket=self._config_reset_socket,
-                )
         except ModbusException as exception_error:
             self._log_error(str(exception_error), error_state=False)
             return False
 
+        for entry in self._pb_call.values():
+            entry[ENTRY_FUNC] = getattr(self._client, entry[ENTRY_NAME])
+
         async with self._lock:
             if not await self.hass.async_add_executor_job(self._pymodbus_connect):
                 self._log_error("initial connect failed, no retry", error_state=False)
                 return False
 
-        self._call_type[CALL_TYPE_COIL][ENTRY_FUNC] = self._client.read_coils
-        self._call_type[CALL_TYPE_DISCRETE][
-            ENTRY_FUNC
-        ] = self._client.read_discrete_inputs
-        self._call_type[CALL_TYPE_REGISTER_HOLDING][
-            ENTRY_FUNC
-        ] = self._client.read_holding_registers
-        self._call_type[CALL_TYPE_REGISTER_INPUT][
-            ENTRY_FUNC
-        ] = self._client.read_input_registers
-        self._call_type[CALL_TYPE_WRITE_COIL][ENTRY_FUNC] = self._client.write_coil
-        self._call_type[CALL_TYPE_WRITE_COILS][ENTRY_FUNC] = self._client.write_coils
-        self._call_type[CALL_TYPE_WRITE_REGISTER][
-            ENTRY_FUNC
-        ] = self._client.write_register
-        self._call_type[CALL_TYPE_WRITE_REGISTERS][
-            ENTRY_FUNC
-        ] = self._client.write_registers
-
         # Start counting down to allow modbus requests.
         if self._config_delay:
             self._async_cancel_listener = async_call_later(
@@ -323,11 +295,11 @@ class ModbusHub:
         """Call sync. pymodbus."""
         kwargs = {"unit": unit} if unit else {}
         try:
-            result = self._call_type[use_call][ENTRY_FUNC](address, value, **kwargs)
+            result = self._pb_call[use_call][ENTRY_FUNC](address, value, **kwargs)
         except ModbusException as exception_error:
             self._log_error(str(exception_error))
             return None
-        if not hasattr(result, self._call_type[use_call][ENTRY_ATTR]):
+        if not hasattr(result, self._pb_call[use_call][ENTRY_ATTR]):
             self._log_error(str(result))
             return None
         self._in_error = False