From 982e75a15fa5ce33b272fde35906801933893996 Mon Sep 17 00:00:00 2001 From: Jan Bouwhuis Date: Mon, 21 Nov 2022 09:20:11 +0100 Subject: [PATCH] Improve type hints MQTT light schema template (#82211) * Improve type hints MQTT light schema template * A few improvements * Follow up comments --- .../components/mqtt/light/schema_template.py | 169 +++++++++--------- 1 file changed, 86 insertions(+), 83 deletions(-) diff --git a/homeassistant/components/mqtt/light/schema_template.py b/homeassistant/components/mqtt/light/schema_template.py index 1085bcb5ef12..7e3c8a0a751c 100644 --- a/homeassistant/components/mqtt/light/schema_template.py +++ b/homeassistant/components/mqtt/light/schema_template.py @@ -1,7 +1,9 @@ """Support for MQTT Template lights.""" from __future__ import annotations +from collections.abc import Callable import logging +from typing import Any import voluptuous as vol @@ -30,7 +32,7 @@ from homeassistant.core import HomeAssistant, callback import homeassistant.helpers.config_validation as cv from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.restore_state import RestoreEntity -from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType +from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType, TemplateVarsType import homeassistant.util.color as color_util from .. import subscription @@ -45,7 +47,13 @@ from ..const import ( ) from ..debug_info import log_messages from ..mixins import MQTT_ENTITY_COMMON_SCHEMA, MqttEntity -from ..models import MqttValueTemplate +from ..models import ( + MqttCommandTemplate, + MqttValueTemplate, + PublishPayloadType, + ReceiveMessage, + ReceivePayloadType, +) from ..util import get_mqtt_data from .schema import MQTT_LIGHT_SCHEMA_SCHEMA from .schema_basic import MQTT_LIGHT_ATTRIBUTES_BLOCKED @@ -70,6 +78,17 @@ CONF_MIN_MIREDS = "min_mireds" CONF_RED_TEMPLATE = "red_template" CONF_WHITE_VALUE_TEMPLATE = "white_value_template" +COMMAND_TEMPLATES = (CONF_COMMAND_ON_TEMPLATE, CONF_COMMAND_OFF_TEMPLATE) +VALUE_TEMPLATES = ( + CONF_BLUE_TEMPLATE, + CONF_BRIGHTNESS_TEMPLATE, + CONF_COLOR_TEMP_TEMPLATE, + CONF_EFFECT_TEMPLATE, + CONF_GREEN_TEMPLATE, + CONF_RED_TEMPLATE, + CONF_STATE_TEMPLATE, +) + _PLATFORM_SCHEMA_BASE = ( MQTT_RW_SCHEMA.extend( { @@ -127,24 +146,30 @@ class MqttLightTemplate(MqttEntity, LightEntity, RestoreEntity): _entity_id_format = ENTITY_ID_FORMAT _attributes_extra_blocked = MQTT_LIGHT_ATTRIBUTES_BLOCKED + _optimistic: bool + _command_templates: dict[ + str, Callable[[PublishPayloadType, TemplateVarsType], PublishPayloadType] + ] + _value_templates: dict[str, Callable[[ReceivePayloadType], ReceivePayloadType]] + _fixed_color_mode: ColorMode | str | None + _topics: dict[str, str | None] - def __init__(self, hass, config, config_entry, discovery_data): + def __init__( + self, + hass: HomeAssistant, + config: ConfigType, + config_entry: ConfigEntry, + discovery_data: DiscoveryInfoType | None, + ) -> None: """Initialize a MQTT Template light.""" - self._topics = None - self._templates = None - self._optimistic = False - - # features - self._fixed_color_mode = None - MqttEntity.__init__(self, hass, config, config_entry, discovery_data) @staticmethod - def config_schema(): + def config_schema() -> vol.Schema: """Return the config schema.""" return DISCOVERY_SCHEMA_TEMPLATE - def _setup_from_config(self, config): + def _setup_from_config(self, config: ConfigType) -> None: """(Re)Setup the entity.""" self._attr_max_mireds = config.get(CONF_MAX_MIREDS, super().max_mireds) self._attr_min_mireds = config.get(CONF_MIN_MIREDS, super().min_mireds) @@ -153,41 +178,37 @@ class MqttLightTemplate(MqttEntity, LightEntity, RestoreEntity): self._topics = { key: config.get(key) for key in (CONF_STATE_TOPIC, CONF_COMMAND_TOPIC) } - self._templates = { - key: config.get(key) - for key in ( - CONF_BLUE_TEMPLATE, - CONF_BRIGHTNESS_TEMPLATE, - CONF_COLOR_TEMP_TEMPLATE, - CONF_COMMAND_OFF_TEMPLATE, - CONF_COMMAND_ON_TEMPLATE, - CONF_EFFECT_TEMPLATE, - CONF_GREEN_TEMPLATE, - CONF_RED_TEMPLATE, - CONF_STATE_TEMPLATE, - ) + self._command_templates = { + key: MqttCommandTemplate(config[key], entity=self).async_render + for key in COMMAND_TEMPLATES } - optimistic = config[CONF_OPTIMISTIC] + self._value_templates = { + key: MqttValueTemplate( + config.get(key), entity=self + ).async_render_with_possible_json_value + for key in VALUE_TEMPLATES + } + optimistic: bool = config[CONF_OPTIMISTIC] self._optimistic = ( optimistic or self._topics[CONF_STATE_TOPIC] is None - or self._templates[CONF_STATE_TEMPLATE] is None + or CONF_STATE_TEMPLATE not in self._config ) color_modes = {ColorMode.ONOFF} - if self._templates[CONF_BRIGHTNESS_TEMPLATE] is not None: + if CONF_BRIGHTNESS_TEMPLATE in config: color_modes.add(ColorMode.BRIGHTNESS) - if self._templates[CONF_COLOR_TEMP_TEMPLATE] is not None: + if CONF_COLOR_TEMP_TEMPLATE in config: color_modes.add(ColorMode.COLOR_TEMP) if ( - self._templates[CONF_RED_TEMPLATE] is not None - and self._templates[CONF_GREEN_TEMPLATE] is not None - and self._templates[CONF_BLUE_TEMPLATE] is not None + CONF_RED_TEMPLATE in config + and CONF_GREEN_TEMPLATE in config + and CONF_BLUE_TEMPLATE in config ): color_modes.add(ColorMode.HS) self._attr_supported_color_modes = filter_supported_color_modes(color_modes) self._fixed_color_mode = None - if len(self.supported_color_modes) == 1: + if self.supported_color_modes and len(self.supported_color_modes) == 1: self._fixed_color_mode = next(iter(self.supported_color_modes)) self._attr_color_mode = self._fixed_color_mode @@ -196,26 +217,21 @@ class MqttLightTemplate(MqttEntity, LightEntity, RestoreEntity): features = features | LightEntityFeature.EFFECT self._attr_supported_features = features - def _update_color_mode(self): + def _update_color_mode(self) -> None: """Update the color_mode attribute.""" if self._fixed_color_mode: return # Support for ct + hs, prioritize hs self._attr_color_mode = ColorMode.HS if self.hs_color else ColorMode.COLOR_TEMP - def _prepare_subscribe_topics(self): + def _prepare_subscribe_topics(self) -> None: """(Re)Subscribe to topics.""" - for tpl in self._templates.values(): - if tpl is not None: - tpl = MqttValueTemplate(tpl, entity=self) @callback @log_messages(self.hass, self.entity_id) - def state_received(msg): + def state_received(msg: ReceiveMessage) -> None: """Handle new MQTT messages.""" - state = self._templates[ - CONF_STATE_TEMPLATE - ].async_render_with_possible_json_value(msg.payload) + state = self._value_templates[CONF_STATE_TEMPLATE](msg.payload) if state == STATE_ON: self._attr_is_on = True elif state == STATE_OFF: @@ -225,21 +241,19 @@ class MqttLightTemplate(MqttEntity, LightEntity, RestoreEntity): else: _LOGGER.warning("Invalid state value received") - if self._templates[CONF_BRIGHTNESS_TEMPLATE] is not None: + if CONF_BRIGHTNESS_TEMPLATE in self._config: try: self._attr_brightness = int( - self._templates[ - CONF_BRIGHTNESS_TEMPLATE - ].async_render_with_possible_json_value(msg.payload) + self._value_templates[CONF_BRIGHTNESS_TEMPLATE](msg.payload) ) except ValueError: _LOGGER.warning("Invalid brightness value received") - if self._templates[CONF_COLOR_TEMP_TEMPLATE] is not None: + if CONF_COLOR_TEMP_TEMPLATE in self._config: try: - color_temp = self._templates[ - CONF_COLOR_TEMP_TEMPLATE - ].async_render_with_possible_json_value(msg.payload) + color_temp = self._value_templates[CONF_COLOR_TEMP_TEMPLATE]( + msg.payload + ) self._attr_color_temp = ( int(color_temp) if color_temp != "None" else None ) @@ -247,20 +261,14 @@ class MqttLightTemplate(MqttEntity, LightEntity, RestoreEntity): _LOGGER.warning("Invalid color temperature value received") if ( - self._templates[CONF_RED_TEMPLATE] is not None - and self._templates[CONF_GREEN_TEMPLATE] is not None - and self._templates[CONF_BLUE_TEMPLATE] is not None + CONF_RED_TEMPLATE in self._config + and CONF_GREEN_TEMPLATE in self._config + and CONF_BLUE_TEMPLATE in self._config ): try: - red = self._templates[ - CONF_RED_TEMPLATE - ].async_render_with_possible_json_value(msg.payload) - green = self._templates[ - CONF_GREEN_TEMPLATE - ].async_render_with_possible_json_value(msg.payload) - blue = self._templates[ - CONF_BLUE_TEMPLATE - ].async_render_with_possible_json_value(msg.payload) + red = self._value_templates[CONF_RED_TEMPLATE](msg.payload) + green = self._value_templates[CONF_GREEN_TEMPLATE](msg.payload) + blue = self._value_templates[CONF_BLUE_TEMPLATE](msg.payload) if red == "None" and green == "None" and blue == "None": self._attr_hs_color = None else: @@ -271,12 +279,11 @@ class MqttLightTemplate(MqttEntity, LightEntity, RestoreEntity): except ValueError: _LOGGER.warning("Invalid color value received") - if self._templates[CONF_EFFECT_TEMPLATE] is not None: - effect = self._templates[ - CONF_EFFECT_TEMPLATE - ].async_render_with_possible_json_value(msg.payload) - - if effect in self._config.get(CONF_EFFECT_LIST): + if CONF_EFFECT_TEMPLATE in self._config: + effect = str(self._value_templates[CONF_EFFECT_TEMPLATE](msg.payload)) + if ( + effect_list := self._config[CONF_EFFECT_LIST] + ) and effect in effect_list: self._attr_effect = effect else: _LOGGER.warning("Unsupported effect value received") @@ -297,7 +304,7 @@ class MqttLightTemplate(MqttEntity, LightEntity, RestoreEntity): }, ) - async def _subscribe_topics(self): + async def _subscribe_topics(self) -> None: """(Re)Subscribe to topics.""" await subscription.async_subscribe_topics(self.hass, self._sub_state) @@ -315,16 +322,16 @@ class MqttLightTemplate(MqttEntity, LightEntity, RestoreEntity): self._attr_effect = last_state.attributes.get(ATTR_EFFECT) @property - def assumed_state(self): + def assumed_state(self) -> bool: """Return True if unable to access real state of the entity.""" return self._optimistic - async def async_turn_on(self, **kwargs): + async def async_turn_on(self, **kwargs: Any) -> None: """Turn the entity on. This method is a coroutine. """ - values = {"state": True} + values: dict[str, Any] = {"state": True} if self._optimistic: self._attr_is_on = True @@ -347,7 +354,7 @@ class MqttLightTemplate(MqttEntity, LightEntity, RestoreEntity): # If there's a brightness topic set, we don't want to scale the RGB # values given using the brightness. - if self._templates[CONF_BRIGHTNESS_TEMPLATE] is not None: + if CONF_BRIGHTNESS_TEMPLATE in self._config: brightness = 255 else: brightness = kwargs.get( @@ -381,10 +388,8 @@ class MqttLightTemplate(MqttEntity, LightEntity, RestoreEntity): values["transition"] = kwargs[ATTR_TRANSITION] await self.async_publish( - self._topics[CONF_COMMAND_TOPIC], - self._templates[CONF_COMMAND_ON_TEMPLATE].async_render( - parse_result=False, **values - ), + str(self._topics[CONF_COMMAND_TOPIC]), + self._command_templates[CONF_COMMAND_ON_TEMPLATE](None, values), self._config[CONF_QOS], self._config[CONF_RETAIN], self._config[CONF_ENCODING], @@ -393,12 +398,12 @@ class MqttLightTemplate(MqttEntity, LightEntity, RestoreEntity): if self._optimistic: self.async_write_ha_state() - async def async_turn_off(self, **kwargs): + async def async_turn_off(self, **kwargs: Any) -> None: """Turn the entity off. This method is a coroutine. """ - values = {"state": False} + values: dict[str, Any] = {"state": False} if self._optimistic: self._attr_is_on = False @@ -406,10 +411,8 @@ class MqttLightTemplate(MqttEntity, LightEntity, RestoreEntity): values["transition"] = kwargs[ATTR_TRANSITION] await self.async_publish( - self._topics[CONF_COMMAND_TOPIC], - self._templates[CONF_COMMAND_OFF_TEMPLATE].async_render( - parse_result=False, **values - ), + str(self._topics[CONF_COMMAND_TOPIC]), + self._command_templates[CONF_COMMAND_OFF_TEMPLATE](None, values), self._config[CONF_QOS], self._config[CONF_RETAIN], self._config[CONF_ENCODING],