Reorg device automation (#26880)

* async_trigger -> async_attach_trigger

* Reorg device automations

* Update docstrings

* Fix types

* Fix extending schemas

homeassistant/components/device_automation/__init__.py

@@ -1,16 +1,12 @@
 """Helpers for device automations."""
 import asyncio
 import logging
-from typing import Callable, cast
 
 import voluptuous as vol
 
+from homeassistant.const import CONF_PLATFORM, CONF_DOMAIN, CONF_DEVICE_ID
 from homeassistant.components import websocket_api
-from homeassistant.const import CONF_DOMAIN
-from homeassistant.core import split_entity_id, HomeAssistant
-import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity_registry import async_entries_for_device
-from homeassistant.helpers.typing import ConfigType
 from homeassistant.loader import async_get_integration, IntegrationNotFound
 
 DOMAIN = "device_automation"
@@ -18,6 +14,21 @@ DOMAIN = "device_automation"
 _LOGGER = logging.getLogger(__name__)
 
 
+TRIGGER_BASE_SCHEMA = vol.Schema(
+    {
+        vol.Required(CONF_PLATFORM): "device",
+        vol.Required(CONF_DOMAIN): str,
+        vol.Required(CONF_DEVICE_ID): str,
+    }
+)
+
+TYPES = {
+    "trigger": ("device_trigger", "async_get_triggers"),
+    "condition": ("device_condition", "async_get_conditions"),
+    "action": ("device_action", "async_get_actions"),
+}
+
+
 async def async_setup(hass, config):
     """Set up device automation."""
     hass.components.websocket_api.async_register_command(
@@ -32,21 +43,9 @@ async def async_setup(hass, config):
     return True
 
 
-async def async_device_condition_from_config(
-    hass: HomeAssistant, config: ConfigType, config_validation: bool = True
-) -> Callable[..., bool]:
-    """Wrap action method with state based condition."""
-    if config_validation:
-        config = cv.DEVICE_CONDITION_SCHEMA(config)
-    integration = await async_get_integration(hass, config[CONF_DOMAIN])
-    platform = integration.get_platform("device_automation")
-    return cast(
-        Callable[..., bool],
-        platform.async_condition_from_config(config, config_validation),  # type: ignore
-    )
-
-
-async def _async_get_device_automations_from_domain(hass, domain, fname, device_id):
+async def _async_get_device_automations_from_domain(
+    hass, domain, automation_type, device_id
+):
     """List device automations."""
     integration = None
     try:
@@ -55,17 +54,18 @@ async def _async_get_device_automations_from_domain(hass, domain, fname, device_
         _LOGGER.warning("Integration %s not found", domain)
         return None
 
+    platform_name, function_name = TYPES[automation_type]
+
     try:
-        platform = integration.get_platform("device_automation")
+        platform = integration.get_platform(platform_name)
     except ImportError:
         # The domain does not have device automations
         return None
 
-    if hasattr(platform, fname):
-        return await getattr(platform, fname)(hass, device_id)
+    return await getattr(platform, function_name)(hass, device_id)
 
 
-async def _async_get_device_automations(hass, fname, device_id):
+async def _async_get_device_automations(hass, automation_type, device_id):
     """List device automations."""
     device_registry, entity_registry = await asyncio.gather(
         hass.helpers.device_registry.async_get_registry(),
@@ -79,13 +79,15 @@ async def _async_get_device_automations(hass, fname, device_id):
         config_entry = hass.config_entries.async_get_entry(entry_id)
         domains.add(config_entry.domain)
 
-    entities = async_entries_for_device(entity_registry, device_id)
-    for entity in entities:
-        domains.add(split_entity_id(entity.entity_id)[0])
+    entity_entries = async_entries_for_device(entity_registry, device_id)
+    for entity_entry in entity_entries:
+        domains.add(entity_entry.domain)
 
     device_automations = await asyncio.gather(
         *(
-            _async_get_device_automations_from_domain(hass, domain, fname, device_id)
+            _async_get_device_automations_from_domain(
+                hass, domain, automation_type, device_id
+            )
             for domain in domains
         )
     )
@@ -106,7 +108,7 @@ async def _async_get_device_automations(hass, fname, device_id):
 async def websocket_device_automation_list_actions(hass, connection, msg):
     """Handle request for device actions."""
     device_id = msg["device_id"]
-    actions = await _async_get_device_automations(hass, "async_get_actions", device_id)
+    actions = await _async_get_device_automations(hass, "action", device_id)
     connection.send_result(msg["id"], actions)
 
 
@@ -120,9 +122,7 @@ async def websocket_device_automation_list_actions(hass, connection, msg):
 async def websocket_device_automation_list_conditions(hass, connection, msg):
     """Handle request for device conditions."""
     device_id = msg["device_id"]
-    conditions = await _async_get_device_automations(
-        hass, "async_get_conditions", device_id
-    )
+    conditions = await _async_get_device_automations(hass, "condition", device_id)
     connection.send_result(msg["id"], conditions)
 
 
@@ -136,7 +136,5 @@ async def websocket_device_automation_list_conditions(hass, connection, msg):
 async def websocket_device_automation_list_triggers(hass, connection, msg):
     """Handle request for device triggers."""
     device_id = msg["device_id"]
-    triggers = await _async_get_device_automations(
-        hass, "async_get_triggers", device_id
-    )
+    triggers = await _async_get_device_automations(hass, "trigger", device_id)
     connection.send_result(msg["id"], triggers)

homeassistant/components/device_automation/toggle_entity.py

@@ -1,7 +1,9 @@
 """Device automation helpers for toggle entity."""
+from typing import List
 import voluptuous as vol
 
-import homeassistant.components.automation.state as state
+from homeassistant.core import Context, HomeAssistant, CALLBACK_TYPE
+from homeassistant.components.automation import state, AutomationActionType
 from homeassistant.components.device_automation.const import (
     CONF_IS_OFF,
     CONF_IS_ON,
@@ -11,17 +13,11 @@ from homeassistant.components.device_automation.const import (
     CONF_TURNED_OFF,
     CONF_TURNED_ON,
 )
-from homeassistant.core import split_entity_id
-from homeassistant.const import (
-    CONF_CONDITION,
-    CONF_DEVICE_ID,
-    CONF_DOMAIN,
-    CONF_ENTITY_ID,
-    CONF_PLATFORM,
-    CONF_TYPE,
-)
+from homeassistant.const import CONF_CONDITION, CONF_ENTITY_ID, CONF_PLATFORM, CONF_TYPE
 from homeassistant.helpers.entity_registry import async_entries_for_device
 from homeassistant.helpers import condition, config_validation as cv, service
+from homeassistant.helpers.typing import ConfigType, TemplateVarsType
+from . import TRIGGER_BASE_SCHEMA
 
 ENTITY_ACTIONS = [
     {
@@ -64,41 +60,35 @@ ENTITY_TRIGGERS = [
     },
 ]
 
-ACTION_SCHEMA = vol.Schema(
+ACTION_SCHEMA = cv.DEVICE_ACTION_BASE_SCHEMA.extend(
     {
-        vol.Required(CONF_DEVICE_ID): str,
-        vol.Required(CONF_DOMAIN): str,
         vol.Required(CONF_ENTITY_ID): cv.entity_id,
         vol.Required(CONF_TYPE): vol.In([CONF_TOGGLE, CONF_TURN_OFF, CONF_TURN_ON]),
     }
 )
 
-CONDITION_SCHEMA = vol.Schema(
+CONDITION_SCHEMA = cv.DEVICE_CONDITION_BASE_SCHEMA.extend(
     {
-        vol.Required(CONF_CONDITION): "device",
-        vol.Required(CONF_DEVICE_ID): str,
-        vol.Required(CONF_DOMAIN): str,
         vol.Required(CONF_ENTITY_ID): cv.entity_id,
         vol.Required(CONF_TYPE): vol.In([CONF_IS_OFF, CONF_IS_ON]),
     }
 )
 
-TRIGGER_SCHEMA = vol.Schema(
+TRIGGER_SCHEMA = TRIGGER_BASE_SCHEMA.extend(
     {
-        vol.Required(CONF_PLATFORM): "device",
-        vol.Required(CONF_DEVICE_ID): str,
-        vol.Required(CONF_DOMAIN): str,
         vol.Required(CONF_ENTITY_ID): cv.entity_id,
         vol.Required(CONF_TYPE): vol.In([CONF_TURNED_OFF, CONF_TURNED_ON]),
     }
 )
 
 
-def _is_domain(entity, domain):
-    return split_entity_id(entity.entity_id)[0] == domain
-
-
-async def async_call_action_from_config(hass, config, variables, context, domain):
+async def async_call_action_from_config(
+    hass: HomeAssistant,
+    config: ConfigType,
+    variables: TemplateVarsType,
+    context: Context,
+    domain: str,
+):
     """Change state based on configuration."""
     config = ACTION_SCHEMA(config)
     action_type = config[CONF_TYPE]
@@ -119,7 +109,9 @@ async def async_call_action_from_config(hass, config, variables, context, domain
     )
 
 
-def async_condition_from_config(config, config_validation):
+def async_condition_from_config(
+    config: ConfigType, config_validation: bool
+) -> condition.ConditionCheckerType:
     """Evaluate state based on configuration."""
     condition_type = config[CONF_TYPE]
     if condition_type == CONF_IS_ON:
@@ -135,7 +127,12 @@ def async_condition_from_config(config, config_validation):
     return condition.state_from_config(state_config, config_validation)
 
 
-async def async_attach_trigger(hass, config, action, automation_info):
+async def async_attach_trigger(
+    hass: HomeAssistant,
+    config: ConfigType,
+    action: AutomationActionType,
+    automation_info: dict,
+) -> CALLBACK_TYPE:
     """Listen for state changes based on configuration."""
     trigger_type = config[CONF_TYPE]
     if trigger_type == CONF_TURNED_ON:
@@ -150,37 +147,56 @@ async def async_attach_trigger(hass, config, action, automation_info):
         state.CONF_TO: to_state,
     }
 
-    return await state.async_trigger(hass, state_config, action, automation_info)
+    return await state.async_attach_trigger(
+        hass, state_config, action, automation_info, platform_type="device"
+    )
 
 
-async def _async_get_automations(hass, device_id, automation_templates, domain):
+async def _async_get_automations(
+    hass: HomeAssistant, device_id: str, automation_templates: List[dict], domain: str
+) -> List[dict]:
     """List device automations."""
     automations = []
     entity_registry = await hass.helpers.entity_registry.async_get_registry()
 
-    entities = async_entries_for_device(entity_registry, device_id)
-    domain_entities = [x for x in entities if _is_domain(x, domain)]
-    for entity in domain_entities:
-        for automation in automation_templates:
-            automation = dict(automation)
-            automation.update(
-                device_id=device_id, entity_id=entity.entity_id, domain=domain
+    entries = [
+        entry
+        for entry in async_entries_for_device(entity_registry, device_id)
+        if entry.domain == domain
+    ]
+
+    for entry in entries:
+        automations.extend(
+            (
+                {
+                    **template,
+                    "device_id": device_id,
+                    "entity_id": entry.entity_id,
+                    "domain": domain,
+                }
+                for template in automation_templates
             )
-            automations.append(automation)
+        )
 
     return automations
 
 
-async def async_get_actions(hass, device_id, domain):
+async def async_get_actions(
+    hass: HomeAssistant, device_id: str, domain: str
+) -> List[dict]:
     """List device actions."""
     return await _async_get_automations(hass, device_id, ENTITY_ACTIONS, domain)
 
 
-async def async_get_conditions(hass, device_id, domain):
+async def async_get_conditions(
+    hass: HomeAssistant, device_id: str, domain: str
+) -> List[dict]:
     """List device conditions."""
     return await _async_get_automations(hass, device_id, ENTITY_CONDITIONS, domain)
 
 
-async def async_get_triggers(hass, device_id, domain):
+async def async_get_triggers(
+    hass: HomeAssistant, device_id: str, domain: str
+) -> List[dict]:
     """List device triggers."""
     return await _async_get_automations(hass, device_id, ENTITY_TRIGGERS, domain)