From 36c0e2416d63e44f3876410fa668676e41347fe1 Mon Sep 17 00:00:00 2001
From: Regev Brody <regevbr@gmail.com>
Date: Sun, 26 Mar 2023 23:01:35 +0300
Subject: [PATCH] add select_schema to select component (#4545)

* add select_schema to select component

* add select_schema to select component

* fix cr
---
 esphome/components/copy/select/__init__.py    | 15 ++++++----
 .../modbus_controller/select/__init__.py      |  5 ++--
 esphome/components/select/__init__.py         | 29 +++++++++++++++++--
 .../components/template/select/__init__.py    |  7 +++--
 esphome/components/tuya/select/__init__.py    | 21 ++++++++------
 5 files changed, 55 insertions(+), 22 deletions(-)

diff --git a/esphome/components/copy/select/__init__.py b/esphome/components/copy/select/__init__.py
index 7d4c1c7705..6254ed03e1 100644
--- a/esphome/components/copy/select/__init__.py
+++ b/esphome/components/copy/select/__init__.py
@@ -14,12 +14,15 @@ from .. import copy_ns
 CopySelect = copy_ns.class_("CopySelect", select.Select, cg.Component)
 
 
-CONFIG_SCHEMA = select.SELECT_SCHEMA.extend(
-    {
-        cv.GenerateID(): cv.declare_id(CopySelect),
-        cv.Required(CONF_SOURCE_ID): cv.use_id(select.Select),
-    }
-).extend(cv.COMPONENT_SCHEMA)
+CONFIG_SCHEMA = (
+    select.select_schema(CopySelect)
+    .extend(
+        {
+            cv.Required(CONF_SOURCE_ID): cv.use_id(select.Select),
+        }
+    )
+    .extend(cv.COMPONENT_SCHEMA)
+)
 
 FINAL_VALIDATE_SCHEMA = cv.All(
     inherit_property_from(CONF_ICON, CONF_SOURCE_ID),
diff --git a/esphome/components/modbus_controller/select/__init__.py b/esphome/components/modbus_controller/select/__init__.py
index f8ef61ddc4..5692fea3e3 100644
--- a/esphome/components/modbus_controller/select/__init__.py
+++ b/esphome/components/modbus_controller/select/__init__.py
@@ -64,9 +64,10 @@ INTEGER_SENSOR_VALUE_TYPE = {
 }
 
 CONFIG_SCHEMA = cv.All(
-    select.SELECT_SCHEMA.extend(cv.COMPONENT_SCHEMA).extend(
+    select.select_schema(ModbusSelect)
+    .extend(cv.COMPONENT_SCHEMA)
+    .extend(
         {
-            cv.GenerateID(): cv.declare_id(ModbusSelect),
             cv.GenerateID(CONF_MODBUS_CONTROLLER_ID): cv.use_id(ModbusController),
             cv.Required(CONF_ADDRESS): cv.positive_int,
             cv.Optional(CONF_VALUE_TYPE, default="U_WORD"): cv.enum(
diff --git a/esphome/components/select/__init__.py b/esphome/components/select/__init__.py
index b505d89c6f..760f7600b7 100644
--- a/esphome/components/select/__init__.py
+++ b/esphome/components/select/__init__.py
@@ -3,6 +3,8 @@ import esphome.config_validation as cv
 from esphome import automation
 from esphome.components import mqtt
 from esphome.const import (
+    CONF_ENTITY_CATEGORY,
+    CONF_ICON,
     CONF_ID,
     CONF_ON_VALUE,
     CONF_OPTION,
@@ -14,6 +16,7 @@ from esphome.const import (
     CONF_INDEX,
 )
 from esphome.core import CORE, coroutine_with_priority
+from esphome.cpp_generator import MockObjClass
 from esphome.cpp_helpers import setup_entity
 
 CODEOWNERS = ["@esphome/core"]
@@ -43,8 +46,6 @@ SELECT_OPERATION_OPTIONS = {
     "LAST": SelectOperation.SELECT_OP_LAST,
 }
 
-icon = cv.icon
-
 
 SELECT_SCHEMA = cv.ENTITY_BASE_SCHEMA.extend(cv.MQTT_COMMAND_COMPONENT_SCHEMA).extend(
     {
@@ -58,6 +59,30 @@ SELECT_SCHEMA = cv.ENTITY_BASE_SCHEMA.extend(cv.MQTT_COMMAND_COMPONENT_SCHEMA).e
     }
 )
 
+_UNDEF = object()
+
+
+def select_schema(
+    class_: MockObjClass = _UNDEF,
+    *,
+    entity_category: str = _UNDEF,
+    icon: str = _UNDEF,
+):
+    schema = cv.Schema({})
+    if class_ is not _UNDEF:
+        schema = schema.extend({cv.GenerateID(): cv.declare_id(class_)})
+    if entity_category is not _UNDEF:
+        schema = schema.extend(
+            {
+                cv.Optional(
+                    CONF_ENTITY_CATEGORY, default=entity_category
+                ): cv.entity_category
+            }
+        )
+    if icon is not _UNDEF:
+        schema = schema.extend({cv.Optional(CONF_ICON, default=icon): cv.icon})
+    return SELECT_SCHEMA.extend(schema)
+
 
 async def setup_select_core_(var, config, *, options: list[str]):
     await setup_entity(var, config)
diff --git a/esphome/components/template/select/__init__.py b/esphome/components/template/select/__init__.py
index 4eba77119d..d116cbb8ae 100644
--- a/esphome/components/template/select/__init__.py
+++ b/esphome/components/template/select/__init__.py
@@ -43,9 +43,9 @@ def validate(config):
 
 
 CONFIG_SCHEMA = cv.All(
-    select.SELECT_SCHEMA.extend(
+    select.select_schema(TemplateSelect)
+    .extend(
         {
-            cv.GenerateID(): cv.declare_id(TemplateSelect),
             cv.Required(CONF_OPTIONS): cv.All(
                 cv.ensure_list(cv.string_strict), cv.Length(min=1)
             ),
@@ -55,7 +55,8 @@ CONFIG_SCHEMA = cv.All(
             cv.Optional(CONF_INITIAL_OPTION): cv.string_strict,
             cv.Optional(CONF_RESTORE_VALUE): cv.boolean,
         }
-    ).extend(cv.polling_component_schema("60s")),
+    )
+    .extend(cv.polling_component_schema("60s")),
     validate,
 )
 
diff --git a/esphome/components/tuya/select/__init__.py b/esphome/components/tuya/select/__init__.py
index 3d65eda301..dc78b2c3db 100644
--- a/esphome/components/tuya/select/__init__.py
+++ b/esphome/components/tuya/select/__init__.py
@@ -25,15 +25,18 @@ def ensure_option_map(value):
     return value
 
 
-CONFIG_SCHEMA = select.SELECT_SCHEMA.extend(
-    {
-        cv.GenerateID(): cv.declare_id(TuyaSelect),
-        cv.GenerateID(CONF_TUYA_ID): cv.use_id(Tuya),
-        cv.Required(CONF_ENUM_DATAPOINT): cv.uint8_t,
-        cv.Required(CONF_OPTIONS): ensure_option_map,
-        cv.Optional(CONF_OPTIMISTIC, default=False): cv.boolean,
-    }
-).extend(cv.COMPONENT_SCHEMA)
+CONFIG_SCHEMA = (
+    select.select_schema(TuyaSelect)
+    .extend(
+        {
+            cv.GenerateID(CONF_TUYA_ID): cv.use_id(Tuya),
+            cv.Required(CONF_ENUM_DATAPOINT): cv.uint8_t,
+            cv.Required(CONF_OPTIONS): ensure_option_map,
+            cv.Optional(CONF_OPTIMISTIC, default=False): cv.boolean,
+        }
+    )
+    .extend(cv.COMPONENT_SCHEMA)
+)
 
 
 async def to_code(config):