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

---
 .../analog_threshold/binary_sensor.py         | 28 +++++++++------
 esphome/components/binary_sensor/__init__.py  | 35 ++++++++-----------
 .../matrix_keypad/binary_sensor/__init__.py   |  3 +-
 esphome/components/tm1637/binary_sensor.py    |  3 +-
 .../tm1638/binary_sensor/__init__.py          |  3 +-
 5 files changed, 34 insertions(+), 38 deletions(-)

diff --git a/esphome/components/analog_threshold/binary_sensor.py b/esphome/components/analog_threshold/binary_sensor.py
index ef4a6044bf..7b964dfae6 100644
--- a/esphome/components/analog_threshold/binary_sensor.py
+++ b/esphome/components/analog_threshold/binary_sensor.py
@@ -15,18 +15,24 @@ AnalogThresholdBinarySensor = analog_threshold_ns.class_(
 CONF_UPPER = "upper"
 CONF_LOWER = "lower"
 
-CONFIG_SCHEMA = binary_sensor.BINARY_SENSOR_SCHEMA.extend(
-    {
-        cv.GenerateID(): cv.declare_id(AnalogThresholdBinarySensor),
-        cv.Required(CONF_SENSOR_ID): cv.use_id(sensor.Sensor),
-        cv.Required(CONF_THRESHOLD): cv.Any(
-            cv.float_,
-            cv.Schema(
-                {cv.Required(CONF_UPPER): cv.float_, cv.Required(CONF_LOWER): cv.float_}
+CONFIG_SCHEMA = (
+    binary_sensor.binary_sensor_schema(AnalogThresholdBinarySensor)
+    .extend(
+        {
+            cv.Required(CONF_SENSOR_ID): cv.use_id(sensor.Sensor),
+            cv.Required(CONF_THRESHOLD): cv.Any(
+                cv.float_,
+                cv.Schema(
+                    {
+                        cv.Required(CONF_UPPER): cv.float_,
+                        cv.Required(CONF_LOWER): cv.float_,
+                    }
+                ),
             ),
-        ),
-    }
-).extend(cv.COMPONENT_SCHEMA)
+        }
+    )
+    .extend(cv.COMPONENT_SCHEMA)
+)
 
 
 async def to_code(config):
diff --git a/esphome/components/binary_sensor/__init__.py b/esphome/components/binary_sensor/__init__.py
index 91f63de9e3..f4a5c95b12 100644
--- a/esphome/components/binary_sensor/__init__.py
+++ b/esphome/components/binary_sensor/__init__.py
@@ -393,28 +393,21 @@ def binary_sensor_schema(
     entity_category: str = _UNDEF,
     device_class: str = _UNDEF,
 ) -> cv.Schema:
-    schema = BINARY_SENSOR_SCHEMA
+    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
-            }
-        )
-    return schema
+        # Not cv.optional
+        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),
+    ]:
+        if default is not _UNDEF:
+            schema[cv.Optional(key, default=default)] = validator
+
+    return BINARY_SENSOR_SCHEMA.extend(schema)
 
 
 async def setup_binary_sensor_core_(var, config):
diff --git a/esphome/components/matrix_keypad/binary_sensor/__init__.py b/esphome/components/matrix_keypad/binary_sensor/__init__.py
index 204db98650..9ad909f60a 100644
--- a/esphome/components/matrix_keypad/binary_sensor/__init__.py
+++ b/esphome/components/matrix_keypad/binary_sensor/__init__.py
@@ -30,9 +30,8 @@ def check_button(obj):
 
 
 CONFIG_SCHEMA = cv.All(
-    binary_sensor.BINARY_SENSOR_SCHEMA.extend(
+    binary_sensor.binary_sensor_schema(MatrixKeypadBinarySensor).extend(
         {
-            cv.GenerateID(): cv.declare_id(MatrixKeypadBinarySensor),
             cv.GenerateID(CONF_KEYPAD_ID): cv.use_id(MatrixKeypad),
             cv.Optional(CONF_ROW): cv.int_,
             cv.Optional(CONF_COL): cv.int_,
diff --git a/esphome/components/tm1637/binary_sensor.py b/esphome/components/tm1637/binary_sensor.py
index 66b5172358..f14b9bd018 100644
--- a/esphome/components/tm1637/binary_sensor.py
+++ b/esphome/components/tm1637/binary_sensor.py
@@ -9,9 +9,8 @@ tm1637_ns = cg.esphome_ns.namespace("tm1637")
 TM1637Display = tm1637_ns.class_("TM1637Display", cg.PollingComponent)
 TM1637Key = tm1637_ns.class_("TM1637Key", binary_sensor.BinarySensor)
 
-CONFIG_SCHEMA = binary_sensor.BINARY_SENSOR_SCHEMA.extend(
+CONFIG_SCHEMA = binary_sensor.binary_sensor_schema(TM1637Key).extend(
     {
-        cv.GenerateID(): cv.declare_id(TM1637Key),
         cv.GenerateID(CONF_TM1637_ID): cv.use_id(TM1637Display),
         cv.Required(CONF_KEY): cv.int_range(min=0, max=15),
     }
diff --git a/esphome/components/tm1638/binary_sensor/__init__.py b/esphome/components/tm1638/binary_sensor/__init__.py
index 7262d9e9e1..6623228555 100644
--- a/esphome/components/tm1638/binary_sensor/__init__.py
+++ b/esphome/components/tm1638/binary_sensor/__init__.py
@@ -6,9 +6,8 @@ from ..display import tm1638_ns, TM1638Component, CONF_TM1638_ID
 
 TM1638Key = tm1638_ns.class_("TM1638Key", binary_sensor.BinarySensor)
 
-CONFIG_SCHEMA = binary_sensor.BINARY_SENSOR_SCHEMA.extend(
+CONFIG_SCHEMA = binary_sensor.binary_sensor_schema(TM1638Key).extend(
     {
-        cv.GenerateID(): cv.declare_id(TM1638Key),
         cv.GenerateID(CONF_TM1638_ID): cv.use_id(TM1638Component),
         cv.Required(CONF_KEY): cv.int_range(min=0, max=15),
     }