[lock] Move to use `lock_schema(..) instead of LOCK_SCHEMA` (#8728)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Keith Burzinski <kbx81x@gmail.com>

esphome/components/copy/lock/__init__.py

@@ -1,7 +1,7 @@
 import esphome.codegen as cg
 from esphome.components import lock
 import esphome.config_validation as cv
-from esphome.const import CONF_ENTITY_CATEGORY, CONF_ICON, CONF_ID, CONF_SOURCE_ID
+from esphome.const import CONF_ENTITY_CATEGORY, CONF_ICON, CONF_SOURCE_ID
 from esphome.core.entity_helpers import inherit_property_from
 
 from .. import copy_ns
@@ -9,12 +9,15 @@ from .. import copy_ns
 CopyLock = copy_ns.class_("CopyLock", lock.Lock, cg.Component)
 
 
-CONFIG_SCHEMA = lock.LOCK_SCHEMA.extend(
+CONFIG_SCHEMA = (
-    {
+    lock.lock_schema(CopyLock)
-        cv.GenerateID(): cv.declare_id(CopyLock),
+    .extend(
-        cv.Required(CONF_SOURCE_ID): cv.use_id(lock.Lock),
+        {
-    }
+            cv.Required(CONF_SOURCE_ID): cv.use_id(lock.Lock),
-).extend(cv.COMPONENT_SCHEMA)
+        }
+    )
+    .extend(cv.COMPONENT_SCHEMA)
+)
 
 FINAL_VALIDATE_SCHEMA = cv.All(
     inherit_property_from(CONF_ICON, CONF_SOURCE_ID),
@@ -23,8 +26,7 @@ FINAL_VALIDATE_SCHEMA = cv.All(
 
 
 async def to_code(config):
-    var = cg.new_Pvariable(config[CONF_ID])
+    var = await lock.new_lock(config)
-    await lock.register_lock(var, config)
     await cg.register_component(var, config)
 
     source = await cg.get_variable(config[CONF_SOURCE_ID])

esphome/components/lock/__init__.py

@@ -4,6 +4,8 @@ import esphome.codegen as cg
 from esphome.components import mqtt, web_server
 import esphome.config_validation as cv
 from esphome.const import (
+    CONF_ENTITY_CATEGORY,
+    CONF_ICON,
     CONF_ID,
     CONF_MQTT_ID,
     CONF_ON_LOCK,
@@ -12,6 +14,7 @@ from esphome.const import (
     CONF_WEB_SERVER,
 )
 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,7 +46,7 @@ LOCK_STATES = {
 
 validate_lock_state = cv.enum(LOCK_STATES, upper=True)
 
-LOCK_SCHEMA = (
+_LOCK_SCHEMA = (
     cv.ENTITY_BASE_SCHEMA.extend(web_server.WEBSERVER_SORTING_SCHEMA)
     .extend(cv.MQTT_COMMAND_COMPONENT_SCHEMA)
     .extend(
@@ -64,7 +67,28 @@ LOCK_SCHEMA = (
 )
 
 
-async def setup_lock_core_(var, config):
+def lock_schema(
+    class_: MockObjClass = cv.UNDEFINED,
+    *,
+    icon: str = cv.UNDEFINED,
+    entity_category: str = cv.UNDEFINED,
+) -> cv.Schema:
+    schema = {}
+
+    if class_ is not cv.UNDEFINED:
+        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),
+    ]:
+        if default is not cv.UNDEFINED:
+            schema[cv.Optional(key, default=default)] = validator
+
+    return _LOCK_SCHEMA.extend(schema)
+
+
+async def _setup_lock_core(var, config):
     await setup_entity(var, config)
 
     for conf in config.get(CONF_ON_LOCK, []):
@@ -86,7 +110,13 @@ async def register_lock(var, config):
     if not CORE.has_id(config[CONF_ID]):
         var = cg.Pvariable(config[CONF_ID], var)
     cg.add(cg.App.register_lock(var))
-    await setup_lock_core_(var, config)
+    await _setup_lock_core(var, config)
+
+
+async def new_lock(config, *args):
+    var = cg.new_Pvariable(config[CONF_ID], *args)
+    await register_lock(var, config)
+    return var
 
 
 LOCK_ACTION_SCHEMA = maybe_simple_id(

esphome/components/output/lock/__init__.py

@@ -1,24 +1,26 @@
 import esphome.codegen as cg
 from esphome.components import lock, output
 import esphome.config_validation as cv
-from esphome.const import CONF_ID, CONF_OUTPUT
+from esphome.const import CONF_OUTPUT
 
 from .. import output_ns
 
 OutputLock = output_ns.class_("OutputLock", lock.Lock, cg.Component)
 
-CONFIG_SCHEMA = lock.LOCK_SCHEMA.extend(
+CONFIG_SCHEMA = (
-    {
+    lock.lock_schema(OutputLock)
-        cv.GenerateID(): cv.declare_id(OutputLock),
+    .extend(
-        cv.Required(CONF_OUTPUT): cv.use_id(output.BinaryOutput),
+        {
-    }
+            cv.Required(CONF_OUTPUT): cv.use_id(output.BinaryOutput),
-).extend(cv.COMPONENT_SCHEMA)
+        }
+    )
+    .extend(cv.COMPONENT_SCHEMA)
+)
 
 
 async def to_code(config):
-    var = cg.new_Pvariable(config[CONF_ID])
+    var = await lock.new_lock(config)
     await cg.register_component(var, config)
-    await lock.register_lock(var, config)
 
     output_ = await cg.get_variable(config[CONF_OUTPUT])
     cg.add(var.set_output(output_))

esphome/components/template/lock/__init__.py

@@ -36,9 +36,9 @@ def validate(config):
 
 
 CONFIG_SCHEMA = cv.All(
-    lock.LOCK_SCHEMA.extend(
+    lock.lock_schema(TemplateLock)
+    .extend(
         {
-            cv.GenerateID(): cv.declare_id(TemplateLock),
             cv.Optional(CONF_LAMBDA): cv.returning_lambda,
             cv.Optional(CONF_OPTIMISTIC, default=False): cv.boolean,
             cv.Optional(CONF_ASSUMED_STATE, default=False): cv.boolean,
@@ -48,15 +48,15 @@ CONFIG_SCHEMA = cv.All(
             cv.Optional(CONF_LOCK_ACTION): automation.validate_automation(single=True),
             cv.Optional(CONF_OPEN_ACTION): automation.validate_automation(single=True),
         }
-    ).extend(cv.COMPONENT_SCHEMA),
+    )
+    .extend(cv.COMPONENT_SCHEMA),
     validate,
 )
 
 
 async def to_code(config):
-    var = cg.new_Pvariable(config[CONF_ID])
+    var = await lock.new_lock(config)
     await cg.register_component(var, config)
-    await lock.register_lock(var, config)
 
     if CONF_LAMBDA in config:
         template_ = await cg.process_lambda(