Improve type hints MQTT light schema template (#82211)

* Improve type hints MQTT light schema template

* A few improvements

* Follow up comments
This commit is contained in:
Jan Bouwhuis 2022-11-21 09:20:11 +01:00 committed by GitHub
parent aa3bd78f7e
commit 982e75a15f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 86 additions and 83 deletions

View File

@ -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],