From 40e2832e67fe24c0d581f387cfacf1a5fd781b01 Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <balloob@gmail.com>
Date: Sun, 19 Feb 2023 14:20:13 -0500
Subject: [PATCH] Simplify sensor schema generation (#4462)

* Simplify sensor schema generation

* Mark class not optional

* Fix assignment
---
 esphome/components/sensor/__init__.py | 58 ++++++++-------------------
 1 file changed, 17 insertions(+), 41 deletions(-)

diff --git a/esphome/components/sensor/__init__.py b/esphome/components/sensor/__init__.py
index 54fdfcffdc..d3cb39c2f6 100644
--- a/esphome/components/sensor/__init__.py
+++ b/esphome/components/sensor/__init__.py
@@ -272,48 +272,24 @@ def sensor_schema(
     state_class: str = _UNDEF,
     entity_category: str = _UNDEF,
 ) -> cv.Schema:
-    schema = SENSOR_SCHEMA
+    schema = {}
+
     if class_ is not _UNDEF:
-        schema = schema.extend({cv.GenerateID(): cv.declare_id(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
-            }
-        )
-    if icon is not _UNDEF:
-        schema = schema.extend({cv.Optional(CONF_ICON, default=icon): validate_icon})
-    if accuracy_decimals is not _UNDEF:
-        schema = schema.extend(
-            {
-                cv.Optional(
-                    CONF_ACCURACY_DECIMALS, default=accuracy_decimals
-                ): validate_accuracy_decimals,
-            }
-        )
-    if device_class is not _UNDEF:
-        schema = schema.extend(
-            {
-                cv.Optional(
-                    CONF_DEVICE_CLASS, default=device_class
-                ): validate_device_class
-            }
-        )
-    if state_class is not _UNDEF:
-        schema = schema.extend(
-            {cv.Optional(CONF_STATE_CLASS, default=state_class): validate_state_class}
-        )
-    if entity_category is not _UNDEF:
-        schema = schema.extend(
-            {
-                cv.Optional(
-                    CONF_ENTITY_CATEGORY, default=entity_category
-                ): cv.entity_category
-            }
-        )
-    return schema
+        # Not optional.
+        schema[cv.GenerateID()] = cv.declare_id(class_)
+
+    for key, default, validator in [
+        (CONF_UNIT_OF_MEASUREMENT, unit_of_measurement, validate_unit_of_measurement),
+        (CONF_ICON, icon, validate_icon),
+        (CONF_ACCURACY_DECIMALS, accuracy_decimals, validate_accuracy_decimals),
+        (CONF_DEVICE_CLASS, device_class, validate_device_class),
+        (CONF_STATE_CLASS, state_class, validate_state_class),
+        (CONF_ENTITY_CATEGORY, entity_category, cv.entity_category),
+    ]:
+        if default is not _UNDEF:
+            schema[cv.Optional(key, default=default)] = validator
+
+    return SENSOR_SCHEMA.extend(schema)
 
 
 @FILTER_REGISTRY.register("offset", OffsetFilter, cv.float_)