add SUB_NUMBER macro and schema to number (#4449)

* add SUB_NUMBER macro and schema

* add SUB_NUMBER macro and schema

* add SUB_NUMBER macro and schema

esphome/components/copy/number/__init__.py

@@ -15,12 +15,15 @@ from .. import copy_ns
 CopyNumber = copy_ns.class_("CopyNumber", number.Number, cg.Component)
 
 
-CONFIG_SCHEMA = number.NUMBER_SCHEMA.extend(
+CONFIG_SCHEMA = (
-    {
+    number.number_schema(CopyNumber)
-        cv.GenerateID(): cv.declare_id(CopyNumber),
+    .extend(
-        cv.Required(CONF_SOURCE_ID): cv.use_id(number.Number),
+        {
-    }
+            cv.Required(CONF_SOURCE_ID): cv.use_id(number.Number),
-).extend(cv.COMPONENT_SCHEMA)
+        }
+    )
+    .extend(cv.COMPONENT_SCHEMA)
+)
 
 FINAL_VALIDATE_SCHEMA = cv.All(
     inherit_property_from(CONF_ICON, CONF_SOURCE_ID),

esphome/components/demo/__init__.py

@@ -284,9 +284,10 @@ CONFIG_SCHEMA = cv.Schema(
                 },
             ],
         ): [
-            number.NUMBER_SCHEMA.extend(cv.COMPONENT_SCHEMA).extend(
+            number.number_schema(DemoNumber)
+            .extend(cv.COMPONENT_SCHEMA)
+            .extend(
                 {
-                    cv.GenerateID(): cv.declare_id(DemoNumber),
                     cv.Required(CONF_TYPE): cv.enum(NUMBER_TYPES, int=True),
                     cv.Required(CONF_MIN_VALUE): cv.float_,
                     cv.Required(CONF_MAX_VALUE): cv.float_,

esphome/components/modbus_controller/number/__init__.py

@@ -60,9 +60,10 @@ def validate_modbus_number(config):
 
 
 CONFIG_SCHEMA = cv.All(
-    number.NUMBER_SCHEMA.extend(ModbusItemBaseSchema).extend(
+    number.number_schema(ModbusNumber)
+    .extend(ModbusItemBaseSchema)
+    .extend(
         {
-            cv.GenerateID(): cv.declare_id(ModbusNumber),
             cv.Optional(CONF_REGISTER_TYPE, default="holding"): cv.enum(
                 MODBUS_WRITE_REGISTER_TYPE
             ),

esphome/components/number/__init__.py

@@ -7,7 +7,9 @@ from esphome.const import (
     CONF_ABOVE,
     CONF_BELOW,
     CONF_DEVICE_CLASS,
+    CONF_ENTITY_CATEGORY,
     CONF_ID,
+    CONF_ICON,
     CONF_MODE,
     CONF_ON_VALUE,
     CONF_ON_VALUE_RANGE,
@@ -63,6 +65,7 @@ from esphome.const import (
 )
 from esphome.core import CORE, coroutine_with_priority
 from esphome.cpp_helpers import setup_entity
+from esphome.cpp_generator import MockObjClass
 
 CODEOWNERS = ["@esphome/core"]
 DEVICE_CLASSES = [
@@ -150,8 +153,8 @@ NUMBER_OPERATION_OPTIONS = {
     "TO_MAX": NumberOperation.NUMBER_OP_TO_MAX,
 }
 
-icon = cv.icon
 validate_device_class = cv.one_of(*DEVICE_CLASSES, lower=True, space="_")
+validate_unit_of_measurement = cv.string_strict
 
 NUMBER_SCHEMA = cv.ENTITY_BASE_SCHEMA.extend(cv.MQTT_COMMAND_COMPONENT_SCHEMA).extend(
     {
@@ -170,12 +173,54 @@ NUMBER_SCHEMA = cv.ENTITY_BASE_SCHEMA.extend(cv.MQTT_COMMAND_COMPONENT_SCHEMA).e
             },
             cv.has_at_least_one_key(CONF_ABOVE, CONF_BELOW),
         ),
-        cv.Optional(CONF_UNIT_OF_MEASUREMENT): cv.string_strict,
+        cv.Optional(CONF_UNIT_OF_MEASUREMENT): validate_unit_of_measurement,
         cv.Optional(CONF_MODE, default="AUTO"): cv.enum(NUMBER_MODES, upper=True),
         cv.Optional(CONF_DEVICE_CLASS): validate_device_class,
     }
 )
 
+_UNDEF = object()
+
+
+def number_schema(
+    class_: MockObjClass = _UNDEF,
+    *,
+    icon: str = _UNDEF,
+    entity_category: str = _UNDEF,
+    device_class: str = _UNDEF,
+    unit_of_measurement: str = _UNDEF,
+) -> cv.Schema:
+    schema = NUMBER_SCHEMA
+    if class_ is not _UNDEF:
+        schema = schema.extend({cv.GenerateID(): cv.declare_id(class_)})
+    if icon is not _UNDEF:
+        schema = schema.extend({cv.Optional(CONF_ICON, default=icon): cv.icon})
+    if entity_category is not _UNDEF:
+        schema = schema.extend(
+            {
+                cv.Optional(
+                    CONF_ENTITY_CATEGORY, default=entity_category
+                ): cv.entity_category
+            }
+        )
+    if device_class is not _UNDEF:
+        schema = schema.extend(
+            {
+                cv.Optional(
+                    CONF_DEVICE_CLASS, default=device_class
+                ): validate_device_class
+            }
+        )
+    if unit_of_measurement is not _UNDEF:
+        schema = schema.extend(
+            {
+                cv.Optional(
+                    CONF_UNIT_OF_MEASUREMENT, default=unit_of_measurement
+                ): validate_unit_of_measurement
+            }
+        )
+    return schema
+
 
 async def setup_number_core_(
     var, config, *, min_value: float, max_value: float, step: Optional[float]

esphome/components/number/number.h

@@ -23,6 +23,13 @@ namespace number {
     } \
   }
 
+#define SUB_NUMBER(name) \
+ protected: \
+  number::Number *name##_number_{nullptr}; \
+\
+ public: \
+  void set_##name##_number(number::Number *number) { this->name##_number_ = number; }
+
 class Number;
 
 /** Base-class for all numbers.

esphome/components/template/number/__init__.py

@@ -46,9 +46,9 @@ def validate(config):
 
 
 CONFIG_SCHEMA = cv.All(
-    number.NUMBER_SCHEMA.extend(
+    number.number_schema(TemplateNumber)
+    .extend(
         {
-            cv.GenerateID(): cv.declare_id(TemplateNumber),
             cv.Required(CONF_MAX_VALUE): cv.float_,
             cv.Required(CONF_MIN_VALUE): cv.float_,
             cv.Required(CONF_STEP): cv.positive_float,
@@ -58,7 +58,8 @@ CONFIG_SCHEMA = cv.All(
             cv.Optional(CONF_INITIAL_VALUE): cv.float_,
             cv.Optional(CONF_RESTORE_VALUE): cv.boolean,
         }
-    ).extend(cv.polling_component_schema("60s")),
+    )
+    .extend(cv.polling_component_schema("60s")),
     validate_min_max,
     validate,
 )

esphome/components/tuya/number/__init__.py

@@ -23,16 +23,17 @@ def validate_min_max(config):
 
 
 CONFIG_SCHEMA = cv.All(
-    number.NUMBER_SCHEMA.extend(
+    number.number_schema(TuyaNumber)
+    .extend(
         {
-            cv.GenerateID(): cv.declare_id(TuyaNumber),
             cv.GenerateID(CONF_TUYA_ID): cv.use_id(Tuya),
             cv.Required(CONF_NUMBER_DATAPOINT): cv.uint8_t,
             cv.Required(CONF_MAX_VALUE): cv.float_,
             cv.Required(CONF_MIN_VALUE): cv.float_,
             cv.Required(CONF_STEP): cv.positive_float,
         }
-    ).extend(cv.COMPONENT_SCHEMA),
+    )
+    .extend(cv.COMPONENT_SCHEMA),
     validate_min_max,
 )