From add40c765299a959db67a77047df3022e8aeb48d Mon Sep 17 00:00:00 2001
From: Jesse Hills <3060199+jesserockz@users.noreply.github.com>
Date: Mon, 20 Feb 2023 10:13:35 +1300
Subject: [PATCH] Simplify number_schema function (#4467)

---
 esphome/components/number/__init__.py | 45 ++++++++-------------------
 1 file changed, 13 insertions(+), 32 deletions(-)

diff --git a/esphome/components/number/__init__.py b/esphome/components/number/__init__.py
index a51cf13124..70c53cace3 100644
--- a/esphome/components/number/__init__.py
+++ b/esphome/components/number/__init__.py
@@ -159,7 +159,6 @@ validate_unit_of_measurement = cv.string_strict
 NUMBER_SCHEMA = cv.ENTITY_BASE_SCHEMA.extend(cv.MQTT_COMMAND_COMPONENT_SCHEMA).extend(
     {
         cv.OnlyWith(CONF_MQTT_ID, "mqtt"): cv.declare_id(mqtt.MQTTNumberComponent),
-        cv.GenerateID(): cv.declare_id(Number),
         cv.Optional(CONF_ON_VALUE): automation.validate_automation(
             {
                 cv.GenerateID(CONF_TRIGGER_ID): cv.declare_id(NumberStateTrigger),
@@ -183,43 +182,25 @@ _UNDEF = object()
 
 
 def number_schema(
-    class_: MockObjClass = _UNDEF,
+    class_: MockObjClass,
     *,
     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
+    schema = {cv.GenerateID(): cv.declare_id(class_)}
+
+    for key, default, validator in [
+        (CONF_ICON, icon, cv.icon),
+        (CONF_ENTITY_CATEGORY, entity_category, cv.entity_category),
+        (CONF_DEVICE_CLASS, device_class, validate_device_class),
+        (CONF_UNIT_OF_MEASUREMENT, unit_of_measurement, validate_unit_of_measurement),
+    ]:
+        if default is not _UNDEF:
+            schema[cv.Optional(key, default=default)] = validator
+
+    return NUMBER_SCHEMA.extend(schema)
 
 
 async def setup_number_core_(