Improve MQTT type hints / refactor part 7 - trigger (#81019)

* Improve typing trigger

* Improve hints on templates

* Use new sentinel for template

* Follow-up comment
This commit is contained in:
Jan Bouwhuis 2022-11-08 17:02:06 +01:00 committed by GitHub
parent 2cb58b8189
commit 7d768dc3a6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -1,21 +1,31 @@
"""Offer MQTT listening automation rules.""" """Offer MQTT listening automation rules."""
from __future__ import annotations
from collections.abc import Callable
from contextlib import suppress from contextlib import suppress
import logging import logging
from typing import Any
import voluptuous as vol import voluptuous as vol
from homeassistant.const import CONF_PAYLOAD, CONF_PLATFORM, CONF_VALUE_TEMPLATE from homeassistant.const import CONF_PAYLOAD, CONF_PLATFORM, CONF_VALUE_TEMPLATE
from homeassistant.core import CALLBACK_TYPE, HassJob, HomeAssistant, callback from homeassistant.core import CALLBACK_TYPE, HassJob, HomeAssistant, callback
from homeassistant.helpers import config_validation as cv, template from homeassistant.helpers import config_validation as cv
from homeassistant.helpers.json import json_loads from homeassistant.helpers.json import json_loads
from homeassistant.helpers.trigger import TriggerActionType, TriggerInfo from homeassistant.helpers.template import Template
from homeassistant.helpers.typing import ConfigType from homeassistant.helpers.trigger import TriggerActionType, TriggerData, TriggerInfo
from homeassistant.helpers.typing import ConfigType, TemplateVarsType
from .. import mqtt from .. import mqtt
from .const import CONF_ENCODING, CONF_QOS, CONF_TOPIC, DEFAULT_ENCODING, DEFAULT_QOS from .const import CONF_ENCODING, CONF_QOS, CONF_TOPIC, DEFAULT_ENCODING, DEFAULT_QOS
from .models import (
# mypy: allow-untyped-defs MqttCommandTemplate,
MqttValueTemplate,
PayloadSentinel,
PublishPayloadType,
ReceiveMessage,
ReceivePayloadType,
)
TRIGGER_SCHEMA = cv.TRIGGER_BASE_SCHEMA.extend( TRIGGER_SCHEMA = cv.TRIGGER_BASE_SCHEMA.extend(
{ {
@ -40,43 +50,37 @@ async def async_attach_trigger(
trigger_info: TriggerInfo, trigger_info: TriggerInfo,
) -> CALLBACK_TYPE: ) -> CALLBACK_TYPE:
"""Listen for state changes based on configuration.""" """Listen for state changes based on configuration."""
trigger_data = trigger_info["trigger_data"] trigger_data: TriggerData = trigger_info["trigger_data"]
topic = config[CONF_TOPIC] command_template: Callable[
wanted_payload = config.get(CONF_PAYLOAD) [PublishPayloadType, TemplateVarsType], PublishPayloadType
value_template = config.get(CONF_VALUE_TEMPLATE) ] = MqttCommandTemplate(config.get(CONF_PAYLOAD), hass=hass).async_render
encoding = config[CONF_ENCODING] or None value_template: Callable[[ReceivePayloadType, str], ReceivePayloadType]
qos = config[CONF_QOS] value_template = MqttValueTemplate(
config.get(CONF_VALUE_TEMPLATE), hass=hass
).async_render_with_possible_json_value
encoding: str | None = config[CONF_ENCODING] or None
qos: int = config[CONF_QOS]
job = HassJob(action) job = HassJob(action)
variables = None variables: TemplateVarsType | None = None
if trigger_info: if trigger_info:
variables = trigger_info.get("variables") variables = trigger_info.get("variables")
template.attach(hass, wanted_payload) wanted_payload = command_template(None, variables)
if wanted_payload:
wanted_payload = wanted_payload.async_render(
variables, limited=True, parse_result=False
)
template.attach(hass, topic) topic_template: Template = config[CONF_TOPIC]
if isinstance(topic, template.Template): topic_template.hass = hass
topic = topic.async_render(variables, limited=True, parse_result=False) topic = topic_template.async_render(variables, limited=True, parse_result=False)
topic = mqtt.util.valid_subscribe_topic(topic) mqtt.util.valid_subscribe_topic(topic)
template.attach(hass, value_template)
@callback @callback
def mqtt_automation_listener(mqttmsg): def mqtt_automation_listener(mqttmsg: ReceiveMessage) -> None:
"""Listen for MQTT messages.""" """Listen for MQTT messages."""
payload = mqttmsg.payload if wanted_payload is None or (
(payload := value_template(mqttmsg.payload, PayloadSentinel.DEFAULT))
if value_template is not None: and payload is not PayloadSentinel.DEFAULT
payload = value_template.async_render_with_possible_json_value( and wanted_payload == payload
payload, ):
error_value=None, data: dict[str, Any] = {
)
if wanted_payload is None or wanted_payload == payload:
data = {
**trigger_data, **trigger_data,
"platform": "mqtt", "platform": "mqtt",
"topic": mqttmsg.topic, "topic": mqttmsg.topic,