Modbus_controller: Add custom command. (#2680)

esphome/components/modbus_controller/switch/__init__.py

@@ -3,21 +3,25 @@ import esphome.config_validation as cv
 import esphome.codegen as cg
 
 
-from esphome.const import CONF_ID, CONF_ADDRESS, CONF_LAMBDA, CONF_OFFSET
+from esphome.const import CONF_ID, CONF_ADDRESS
 from .. import (
-    MODBUS_REGISTER_TYPE,
-    SensorItem,
+    add_modbus_base_properties,
     modbus_controller_ns,
-    ModbusController,
+    modbus_calc_properties,
+    validate_modbus_register,
+    ModbusItemBaseSchema,
+    SensorItem,
+    MODBUS_REGISTER_TYPE,
 )
 from ..const import (
     CONF_BITMASK,
-    CONF_BYTE_OFFSET,
     CONF_FORCE_NEW_RANGE,
     CONF_MODBUS_CONTROLLER_ID,
     CONF_REGISTER_TYPE,
+    CONF_WRITE_LAMBDA,
 )
 
+CONF_USE_WRITE_MULTIPLE = "use_write_multiple"
 DEPENDENCIES = ["modbus_controller"]
 CODEOWNERS = ["@martgras"]
 
@@ -26,31 +30,23 @@ ModbusSwitch = modbus_controller_ns.class_(
     "ModbusSwitch", cg.Component, switch.Switch, SensorItem
 )
 
-
 CONFIG_SCHEMA = cv.All(
-    switch.SWITCH_SCHEMA.extend(
+    switch.SWITCH_SCHEMA.extend(cv.COMPONENT_SCHEMA)
+    .extend(ModbusItemBaseSchema)
+    .extend(
         {
             cv.GenerateID(): cv.declare_id(ModbusSwitch),
-            cv.GenerateID(CONF_MODBUS_CONTROLLER_ID): cv.use_id(ModbusController),
-            cv.Required(CONF_REGISTER_TYPE): cv.enum(MODBUS_REGISTER_TYPE),
-            cv.Required(CONF_ADDRESS): cv.positive_int,
-            cv.Optional(CONF_OFFSET, default=0): cv.positive_int,
-            cv.Optional(CONF_BYTE_OFFSET): cv.positive_int,
-            cv.Optional(CONF_BITMASK, default=0x1): cv.hex_uint32_t,
-            cv.Optional(CONF_FORCE_NEW_RANGE, default=False): cv.boolean,
-            cv.Optional(CONF_LAMBDA): cv.returning_lambda,
+            cv.Optional(CONF_REGISTER_TYPE): cv.enum(MODBUS_REGISTER_TYPE),
+            cv.Optional(CONF_USE_WRITE_MULTIPLE, default=False): cv.boolean,
+            cv.Optional(CONF_WRITE_LAMBDA): cv.returning_lambda,
         }
-    ).extend(cv.COMPONENT_SCHEMA),
+    ),
+    validate_modbus_register,
 )
 
 
 async def to_code(config):
-    byte_offset = 0
-    if CONF_OFFSET in config:
-        byte_offset = config[CONF_OFFSET]
-    # A CONF_BYTE_OFFSET setting overrides CONF_OFFSET
-    if CONF_BYTE_OFFSET in config:
-        byte_offset = config[CONF_BYTE_OFFSET]
+    byte_offset, _ = modbus_calc_properties(config)
     var = cg.new_Pvariable(
         config[CONF_ID],
         config[CONF_REGISTER_TYPE],
@@ -63,19 +59,18 @@ async def to_code(config):
     await switch.register_switch(var, config)
 
     paren = await cg.get_variable(config[CONF_MODBUS_CONTROLLER_ID])
-    cg.add(paren.add_sensor_item(var))
     cg.add(var.set_parent(paren))
-    if CONF_LAMBDA in config:
-        publish_template_ = await cg.process_lambda(
-            config[CONF_LAMBDA],
+    cg.add(var.set_use_write_mutiple(config[CONF_USE_WRITE_MULTIPLE]))
+    cg.add(paren.add_sensor_item(var))
+    if CONF_WRITE_LAMBDA in config:
+        template_ = await cg.process_lambda(
+            config[CONF_WRITE_LAMBDA],
             [
                 (ModbusSwitch.operator("ptr"), "item"),
-                (bool, "x"),
-                (
-                    cg.std_vector.template(cg.uint8).operator("const").operator("ref"),
-                    "data",
-                ),
+                (cg.bool_, "x"),
+                (cg.std_vector.template(cg.uint8).operator("ref"), "payload"),
             ],
             return_type=cg.optional.template(bool),
         )
-        cg.add(var.set_template(publish_template_))
+        cg.add(var.set_write_template(template_))
+    await add_modbus_base_properties(var, config, ModbusSwitch, bool, bool)