From 756f71c3822a2f457b42c5b55f8af1151c69c444 Mon Sep 17 00:00:00 2001
From: Jesse Hills <3060199+jesserockz@users.noreply.github.com>
Date: Wed, 16 Mar 2022 13:43:05 +1300
Subject: [PATCH] Allow custom register type for modbus number (#3202)

---
 esphome/components/modbus_controller/number/__init__.py     | 6 ++++++
 esphome/components/modbus_controller/number/modbus_number.h | 6 +++---
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/esphome/components/modbus_controller/number/__init__.py b/esphome/components/modbus_controller/number/__init__.py
index 56ec734315..37a39ff334 100644
--- a/esphome/components/modbus_controller/number/__init__.py
+++ b/esphome/components/modbus_controller/number/__init__.py
@@ -11,6 +11,7 @@ from esphome.const import (
 )
 
 from .. import (
+    MODBUS_WRITE_REGISTER_TYPE,
     add_modbus_base_properties,
     modbus_controller_ns,
     modbus_calc_properties,
@@ -24,6 +25,7 @@ from ..const import (
     CONF_CUSTOM_COMMAND,
     CONF_FORCE_NEW_RANGE,
     CONF_MODBUS_CONTROLLER_ID,
+    CONF_REGISTER_TYPE,
     CONF_SKIP_UPDATES,
     CONF_USE_WRITE_MULTIPLE,
     CONF_VALUE_TYPE,
@@ -61,6 +63,9 @@ CONFIG_SCHEMA = cv.All(
     number.NUMBER_SCHEMA.extend(ModbusItemBaseSchema).extend(
         {
             cv.GenerateID(): cv.declare_id(ModbusNumber),
+            cv.Optional(CONF_REGISTER_TYPE, default="holding"): cv.enum(
+                MODBUS_WRITE_REGISTER_TYPE
+            ),
             cv.Optional(CONF_VALUE_TYPE, default="U_WORD"): cv.enum(SENSOR_VALUE_TYPE),
             cv.Optional(CONF_WRITE_LAMBDA): cv.returning_lambda,
             # 24 bits are the maximum value for fp32 before precison is lost
@@ -81,6 +86,7 @@ async def to_code(config):
     byte_offset, reg_count = modbus_calc_properties(config)
     var = cg.new_Pvariable(
         config[CONF_ID],
+        config[CONF_REGISTER_TYPE],
         config[CONF_ADDRESS],
         byte_offset,
         config[CONF_BITMASK],
diff --git a/esphome/components/modbus_controller/number/modbus_number.h b/esphome/components/modbus_controller/number/modbus_number.h
index 0c525d9c89..aa5c8d1500 100644
--- a/esphome/components/modbus_controller/number/modbus_number.h
+++ b/esphome/components/modbus_controller/number/modbus_number.h
@@ -11,9 +11,9 @@ using value_to_data_t = std::function<float>(float);
 
 class ModbusNumber : public number::Number, public Component, public SensorItem {
  public:
-  ModbusNumber(uint16_t start_address, uint8_t offset, uint32_t bitmask, SensorValueType value_type, int register_count,
-               uint8_t skip_updates, bool force_new_range) {
-    this->register_type = ModbusRegisterType::HOLDING;
+  ModbusNumber(ModbusRegisterType register_type, uint16_t start_address, uint8_t offset, uint32_t bitmask,
+               SensorValueType value_type, int register_count, uint8_t skip_updates, bool force_new_range) {
+    this->register_type = register_type;
     this->start_address = start_address;
     this->offset = offset;
     this->bitmask = bitmask;