From 72108684ea20bc6b8527726a27c5c298919f3e24 Mon Sep 17 00:00:00 2001
From: Martin <25747549+martgras@users.noreply.github.com>
Date: Tue, 26 Oct 2021 11:30:25 +0200
Subject: [PATCH] fix modbus output (#2630)

---
 .../components/modbus_controller/__init__.py   | 15 +++++++++++++++
 .../modbus_controller/number/__init__.py       | 17 +----------------
 .../modbus_controller/output/__init__.py       | 13 ++++++++++++-
 .../modbus_controller/output/modbus_output.h   |  3 ++-
 .../modbus_controller/sensor/__init__.py       | 18 +-----------------
 5 files changed, 31 insertions(+), 35 deletions(-)

diff --git a/esphome/components/modbus_controller/__init__.py b/esphome/components/modbus_controller/__init__.py
index 6b452ea25c..8499cec561 100644
--- a/esphome/components/modbus_controller/__init__.py
+++ b/esphome/components/modbus_controller/__init__.py
@@ -61,6 +61,21 @@ SENSOR_VALUE_TYPE = {
     "FP32_R": SensorValueType.FP32_R,
 }
 
+TYPE_REGISTER_MAP = {
+    "RAW": 1,
+    "U_WORD": 1,
+    "S_WORD": 1,
+    "U_DWORD": 2,
+    "U_DWORD_R": 2,
+    "S_DWORD": 2,
+    "S_DWORD_R": 2,
+    "U_QWORD": 4,
+    "U_QWORDU_R": 4,
+    "S_QWORD": 4,
+    "U_QWORD_R": 4,
+    "FP32": 2,
+    "FP32_R": 2,
+}
 
 MULTI_CONF = True
 
diff --git a/esphome/components/modbus_controller/number/__init__.py b/esphome/components/modbus_controller/number/__init__.py
index c7919bb972..afb69f8798 100644
--- a/esphome/components/modbus_controller/number/__init__.py
+++ b/esphome/components/modbus_controller/number/__init__.py
@@ -17,6 +17,7 @@ from .. import (
     ModbusController,
     SENSOR_VALUE_TYPE,
     SensorItem,
+    TYPE_REGISTER_MAP,
 )
 
 
@@ -39,22 +40,6 @@ ModbusNumber = modbus_controller_ns.class_(
     "ModbusNumber", cg.Component, number.Number, SensorItem
 )
 
-TYPE_REGISTER_MAP = {
-    "RAW": 1,
-    "U_WORD": 1,
-    "S_WORD": 1,
-    "U_DWORD": 2,
-    "U_DWORD_R": 2,
-    "S_DWORD": 2,
-    "S_DWORD_R": 2,
-    "U_QWORD": 4,
-    "U_QWORDU_R": 4,
-    "S_QWORD": 4,
-    "U_QWORD_R": 4,
-    "FP32": 2,
-    "FP32_R": 2,
-}
-
 
 def validate_min_max(config):
     if config[CONF_MAX_VALUE] <= config[CONF_MIN_VALUE]:
diff --git a/esphome/components/modbus_controller/output/__init__.py b/esphome/components/modbus_controller/output/__init__.py
index 9c41fc011c..4aca4db64f 100644
--- a/esphome/components/modbus_controller/output/__init__.py
+++ b/esphome/components/modbus_controller/output/__init__.py
@@ -13,11 +13,13 @@ from .. import (
     SensorItem,
     modbus_controller_ns,
     ModbusController,
+    TYPE_REGISTER_MAP,
 )
 
 from ..const import (
     CONF_BYTE_OFFSET,
     CONF_MODBUS_CONTROLLER_ID,
+    CONF_REGISTER_COUNT,
     CONF_VALUE_TYPE,
     CONF_WRITE_LAMBDA,
 )
@@ -40,6 +42,7 @@ CONFIG_SCHEMA = cv.All(
             cv.Optional(CONF_OFFSET, default=0): cv.positive_int,
             cv.Optional(CONF_BYTE_OFFSET): cv.positive_int,
             cv.Optional(CONF_VALUE_TYPE, default="U_WORD"): cv.enum(SENSOR_VALUE_TYPE),
+            cv.Optional(CONF_REGISTER_COUNT, default=0): cv.positive_int,
             cv.Optional(CONF_WRITE_LAMBDA): cv.returning_lambda,
             cv.Optional(CONF_MULTIPLY, default=1.0): cv.float_,
         }
@@ -54,8 +57,16 @@ async def to_code(config):
     # A CONF_BYTE_OFFSET setting overrides CONF_OFFSET
     if CONF_BYTE_OFFSET in config:
         byte_offset = config[CONF_BYTE_OFFSET]
+    value_type = config[CONF_VALUE_TYPE]
+    reg_count = config[CONF_REGISTER_COUNT]
+    if reg_count == 0:
+        reg_count = TYPE_REGISTER_MAP[value_type]
     var = cg.new_Pvariable(
-        config[CONF_ID], config[CONF_ADDRESS], byte_offset, config[CONF_VALUE_TYPE]
+        config[CONF_ID],
+        config[CONF_ADDRESS],
+        byte_offset,
+        value_type,
+        reg_count,
     )
     await output.register_output(var, config)
     cg.add(var.set_write_multiply(config[CONF_MULTIPLY]))
diff --git a/esphome/components/modbus_controller/output/modbus_output.h b/esphome/components/modbus_controller/output/modbus_output.h
index 053186a321..6e8521854b 100644
--- a/esphome/components/modbus_controller/output/modbus_output.h
+++ b/esphome/components/modbus_controller/output/modbus_output.h
@@ -11,12 +11,13 @@ using value_to_data_t = std::function<float>(float);
 
 class ModbusOutput : public output::FloatOutput, public Component, public SensorItem {
  public:
-  ModbusOutput(uint16_t start_address, uint8_t offset, SensorValueType value_type)
+  ModbusOutput(uint16_t start_address, uint8_t offset, SensorValueType value_type, int register_count)
       : output::FloatOutput(), Component() {
     this->register_type = ModbusRegisterType::HOLDING;
     this->start_address = start_address;
     this->offset = offset;
     this->bitmask = bitmask;
+    this->register_count = register_count;
     this->sensor_value_type = value_type;
     this->skip_updates = 0;
     this->start_address += offset;
diff --git a/esphome/components/modbus_controller/sensor/__init__.py b/esphome/components/modbus_controller/sensor/__init__.py
index 687f3d82fb..82acfe120b 100644
--- a/esphome/components/modbus_controller/sensor/__init__.py
+++ b/esphome/components/modbus_controller/sensor/__init__.py
@@ -9,6 +9,7 @@ from .. import (
     ModbusController,
     MODBUS_REGISTER_TYPE,
     SENSOR_VALUE_TYPE,
+    TYPE_REGISTER_MAP,
 )
 from ..const import (
     CONF_BITMASK,
@@ -29,23 +30,6 @@ ModbusSensor = modbus_controller_ns.class_(
     "ModbusSensor", cg.Component, sensor.Sensor, SensorItem
 )
 
-TYPE_REGISTER_MAP = {
-    "RAW": 1,
-    "U_WORD": 1,
-    "S_WORD": 1,
-    "U_DWORD": 2,
-    "U_DWORD_R": 2,
-    "S_DWORD": 2,
-    "S_DWORD_R": 2,
-    "U_QWORD": 4,
-    "U_QWORDU_R": 4,
-    "S_QWORD": 4,
-    "U_QWORD_R": 4,
-    "FP32": 2,
-    "FP32_R": 2,
-}
-
-
 CONFIG_SCHEMA = cv.All(
     sensor.SENSOR_SCHEMA.extend(
         {