Refactor schema generation in Template integration (#120889)

Co-authored-by: Robert Resch <robert@resch.dev>
This commit is contained in:
dougiteixeira 2024-07-08 04:06:15 -03:00 committed by GitHub
parent e30f315565
commit cf4bd7fd1c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -3,6 +3,7 @@
from __future__ import annotations from __future__ import annotations
from collections.abc import Callable, Coroutine, Mapping from collections.abc import Callable, Coroutine, Mapping
from functools import partial
from typing import Any, cast from typing import Any, cast
import voluptuous as vol import voluptuous as vol
@ -39,13 +40,22 @@ from .const import DOMAIN
from .sensor import async_create_preview_sensor from .sensor import async_create_preview_sensor
from .template_entity import TemplateEntity from .template_entity import TemplateEntity
_SCHEMA_STATE: dict[vol.Marker, Any] = {
vol.Required(CONF_STATE): selector.TemplateSelector(),
}
def generate_schema(domain: str, flow_type: str) -> dict[vol.Marker, Any]:
def generate_schema(domain: str, flow_type: str) -> vol.Schema:
"""Generate schema.""" """Generate schema."""
schema: dict[vol.Marker, Any] = {} schema: dict[vol.Marker, Any] = {}
if domain == Platform.BINARY_SENSOR and flow_type == "config": if flow_type == "config":
schema = { schema = {vol.Required(CONF_NAME): selector.TextSelector()}
if domain == Platform.BINARY_SENSOR:
schema |= _SCHEMA_STATE
if flow_type == "config":
schema |= {
vol.Optional(CONF_DEVICE_CLASS): selector.SelectSelector( vol.Optional(CONF_DEVICE_CLASS): selector.SelectSelector(
selector.SelectSelectorConfig( selector.SelectSelectorConfig(
options=[cls.value for cls in BinarySensorDeviceClass], options=[cls.value for cls in BinarySensorDeviceClass],
@ -53,11 +63,11 @@ def generate_schema(domain: str, flow_type: str) -> dict[vol.Marker, Any]:
translation_key="binary_sensor_device_class", translation_key="binary_sensor_device_class",
sort=True, sort=True,
), ),
) ),
} }
if domain == Platform.SENSOR: if domain == Platform.SENSOR:
schema = { schema |= _SCHEMA_STATE | {
vol.Optional(CONF_UNIT_OF_MEASUREMENT): selector.SelectSelector( vol.Optional(CONF_UNIT_OF_MEASUREMENT): selector.SelectSelector(
selector.SelectSelectorConfig( selector.SelectSelectorConfig(
options=list( options=list(
@ -98,26 +108,12 @@ def generate_schema(domain: str, flow_type: str) -> dict[vol.Marker, Any]:
schema[vol.Optional(CONF_DEVICE_ID)] = selector.DeviceSelector() schema[vol.Optional(CONF_DEVICE_ID)] = selector.DeviceSelector()
return schema return vol.Schema(schema)
def options_schema(domain: str) -> vol.Schema: options_schema = partial(generate_schema, flow_type="options")
"""Generate options schema."""
return vol.Schema(
{vol.Required(CONF_STATE): selector.TemplateSelector()}
| generate_schema(domain, "option"),
)
config_schema = partial(generate_schema, flow_type="config")
def config_schema(domain: str) -> vol.Schema:
"""Generate config schema."""
return vol.Schema(
{
vol.Required(CONF_NAME): selector.TextSelector(),
vol.Required(CONF_STATE): selector.TemplateSelector(),
}
| generate_schema(domain, "config"),
)
async def choose_options_step(options: dict[str, Any]) -> str: async def choose_options_step(options: dict[str, Any]) -> str: