212 lines
7.0 KiB
Python
212 lines
7.0 KiB
Python
"""The Search integration."""
|
|
from collections import defaultdict
|
|
|
|
import voluptuous as vol
|
|
|
|
from homeassistant.components import group, websocket_api
|
|
from homeassistant.components.homeassistant import scene
|
|
from homeassistant.core import HomeAssistant, callback, split_entity_id
|
|
from homeassistant.helpers import device_registry, entity_registry
|
|
|
|
DOMAIN = "search"
|
|
|
|
|
|
async def async_setup(hass: HomeAssistant, config: dict):
|
|
"""Set up the Search component."""
|
|
websocket_api.async_register_command(hass, websocket_search_related)
|
|
return True
|
|
|
|
|
|
@websocket_api.async_response
|
|
@websocket_api.websocket_command(
|
|
{
|
|
vol.Required("type"): "search/related",
|
|
vol.Required("item_type"): vol.In(
|
|
(
|
|
"area",
|
|
"automation",
|
|
"config_entry",
|
|
"device",
|
|
"entity",
|
|
"group",
|
|
"scene",
|
|
"script",
|
|
)
|
|
),
|
|
vol.Required("item_id"): str,
|
|
}
|
|
)
|
|
async def websocket_search_related(hass, connection, msg):
|
|
"""Handle search."""
|
|
searcher = Searcher(
|
|
hass,
|
|
await device_registry.async_get_registry(hass),
|
|
await entity_registry.async_get_registry(hass),
|
|
)
|
|
connection.send_result(
|
|
msg["id"], searcher.async_search(msg["item_type"], msg["item_id"])
|
|
)
|
|
|
|
|
|
class Searcher:
|
|
"""Find related things.
|
|
|
|
Few rules:
|
|
Scenes, scripts, automations and config entries will only be expanded if they are
|
|
the entry point. They won't be expanded if we process them. This is because they
|
|
turn the results into garbage.
|
|
"""
|
|
|
|
# These types won't be further explored. Config entries + Output types.
|
|
DONT_RESOLVE = {"scene", "automation", "script", "group", "config_entry", "area"}
|
|
|
|
def __init__(
|
|
self,
|
|
hass: HomeAssistant,
|
|
device_reg: device_registry.DeviceRegistry,
|
|
entity_reg: entity_registry.EntityRegistry,
|
|
):
|
|
"""Search results."""
|
|
self.hass = hass
|
|
self._device_reg = device_reg
|
|
self._entity_reg = entity_reg
|
|
self.results = defaultdict(set)
|
|
self._to_resolve = set()
|
|
|
|
@callback
|
|
def async_search(self, item_type, item_id):
|
|
"""Find results."""
|
|
self.results[item_type].add(item_id)
|
|
self._to_resolve.add((item_type, item_id))
|
|
|
|
while self._to_resolve:
|
|
search_type, search_id = self._to_resolve.pop()
|
|
getattr(self, f"_resolve_{search_type}")(search_id)
|
|
|
|
# Clean up entity_id items, from the general "entity" type result,
|
|
# that are also found in the specific entity domain type.
|
|
self.results["entity"] -= self.results["script"]
|
|
self.results["entity"] -= self.results["scene"]
|
|
self.results["entity"] -= self.results["automation"]
|
|
self.results["entity"] -= self.results["group"]
|
|
|
|
# Remove entry into graph from search results.
|
|
self.results[item_type].remove(item_id)
|
|
|
|
# Filter out empty sets.
|
|
return {key: val for key, val in self.results.items() if val}
|
|
|
|
@callback
|
|
def _add_or_resolve(self, item_type, item_id):
|
|
"""Add an item to explore."""
|
|
if item_id in self.results[item_type]:
|
|
return
|
|
|
|
self.results[item_type].add(item_id)
|
|
|
|
if item_type not in self.DONT_RESOLVE:
|
|
self._to_resolve.add((item_type, item_id))
|
|
|
|
@callback
|
|
def _resolve_area(self, area_id) -> None:
|
|
"""Resolve an area."""
|
|
for device in device_registry.async_entries_for_area(self._device_reg, area_id):
|
|
self._add_or_resolve("device", device.id)
|
|
|
|
@callback
|
|
def _resolve_device(self, device_id) -> None:
|
|
"""Resolve a device."""
|
|
device_entry = self._device_reg.async_get(device_id)
|
|
# Unlikely entry doesn't exist, but let's guard for bad data.
|
|
if device_entry is not None:
|
|
if device_entry.area_id:
|
|
self._add_or_resolve("area", device_entry.area_id)
|
|
|
|
for config_entry_id in device_entry.config_entries:
|
|
self._add_or_resolve("config_entry", config_entry_id)
|
|
|
|
# We do not resolve device_entry.via_device_id because that
|
|
# device is not related data-wise inside HA.
|
|
|
|
for entity_entry in entity_registry.async_entries_for_device(
|
|
self._entity_reg, device_id
|
|
):
|
|
self._add_or_resolve("entity", entity_entry.entity_id)
|
|
|
|
# Extra: Find automations that reference this device
|
|
|
|
@callback
|
|
def _resolve_entity(self, entity_id) -> None:
|
|
"""Resolve an entity."""
|
|
# Extra: Find automations and scripts that reference this entity.
|
|
|
|
for entity in scene.scenes_with_entity(self.hass, entity_id):
|
|
self._add_or_resolve("entity", entity)
|
|
|
|
for entity in group.groups_with_entity(self.hass, entity_id):
|
|
self._add_or_resolve("entity", entity)
|
|
|
|
# Find devices
|
|
entity_entry = self._entity_reg.async_get(entity_id)
|
|
if entity_entry is not None:
|
|
if entity_entry.device_id:
|
|
self._add_or_resolve("device", entity_entry.device_id)
|
|
|
|
if entity_entry.config_entry_id is not None:
|
|
self._add_or_resolve("config_entry", entity_entry.config_entry_id)
|
|
|
|
domain = split_entity_id(entity_id)[0]
|
|
|
|
if domain in ("scene", "automation", "script", "group"):
|
|
self._add_or_resolve(domain, entity_id)
|
|
|
|
@callback
|
|
def _resolve_automation(self, automation_entity_id) -> None:
|
|
"""Resolve an automation.
|
|
|
|
Will only be called if automation is an entry point.
|
|
"""
|
|
# Extra: Check with automation integration what entities/devices they reference
|
|
|
|
@callback
|
|
def _resolve_script(self, script_entity_id) -> None:
|
|
"""Resolve a script.
|
|
|
|
Will only be called if script is an entry point.
|
|
"""
|
|
# Extra: Check with script integration what entities/devices they reference
|
|
|
|
@callback
|
|
def _resolve_group(self, group_entity_id) -> None:
|
|
"""Resolve a group.
|
|
|
|
Will only be called if group is an entry point.
|
|
"""
|
|
for entity_id in group.get_entity_ids(self.hass, group_entity_id):
|
|
self._add_or_resolve("entity", entity_id)
|
|
|
|
@callback
|
|
def _resolve_scene(self, scene_entity_id) -> None:
|
|
"""Resolve a scene.
|
|
|
|
Will only be called if scene is an entry point.
|
|
"""
|
|
for entity in scene.entities_in_scene(self.hass, scene_entity_id):
|
|
self._add_or_resolve("entity", entity)
|
|
|
|
@callback
|
|
def _resolve_config_entry(self, config_entry_id) -> None:
|
|
"""Resolve a config entry.
|
|
|
|
Will only be called if config entry is an entry point.
|
|
"""
|
|
for device_entry in device_registry.async_entries_for_config_entry(
|
|
self._device_reg, config_entry_id
|
|
):
|
|
self._add_or_resolve("device", device_entry.id)
|
|
|
|
for entity_entry in entity_registry.async_entries_for_config_entry(
|
|
self._entity_reg, config_entry_id
|
|
):
|
|
self._add_or_resolve("entity", entity_entry.entity_id)
|