/
/
/
1"""Searchable listing of the Home Assistant entities that can act as a player control."""
2
3from __future__ import annotations
4
5import asyncio
6import time
7from typing import TYPE_CHECKING, Final, NamedTuple, TypedDict
8
9from music_assistant_models.errors import InvalidDataError, ProviderUnavailableError
10
11from .constants import (
12 CONF_MUTE_CONTROLS,
13 CONF_POWER_CONTROLS,
14 CONF_VOLUME_CONTROLS,
15 CONTROL_DOMAINS,
16)
17from .helpers import get_control_capabilities
18
19if TYPE_CHECKING:
20 from collections.abc import Callable
21
22 from . import HomeAssistantProvider
23
24SEARCH_CONTROL_ENTITIES_LIMIT = 50
25# Ceiling on what a caller may raise the limit to, so a large Home Assistant setup can never
26# be asked for a response big enough to hurt, however the search is called.
27SEARCH_CONTROL_ENTITIES_MAX_LIMIT = 500
28# How long a search may reuse the control entity candidates of an earlier search. Resolving
29# them takes a full state sweep of Home Assistant, so a picker that searches while the user
30# types would otherwise sweep on every keystroke. The candidates carry entity, device and
31# area names plus control capabilities - none of which change often, and none of which is
32# live entity state - so briefly serving a stale listing is harmless. Which entities exist
33# at all is pinned to the entity registry instead of to this window, so an entity that
34# appears or disappears is picked up right away.
35CONTROL_ENTITY_CACHE_TTL = 30
36
37# The entity domains worth inspecting for each control role. A superset is harmless: the
38# authoritative verdict comes from get_control_capabilities on the entity's own state,
39# this only keeps the state sweep of a search away from domains that can never qualify.
40CONTROL_TYPE_DOMAINS: Final[dict[str, tuple[str, ...]]] = {
41 CONF_POWER_CONTROLS: ("media_player", "switch", "input_boolean"),
42 CONF_VOLUME_CONTROLS: ("media_player", "number", "input_number"),
43 CONF_MUTE_CONTROLS: ("media_player", "switch", "input_boolean"),
44}
45
46
47class HassControlEntity(TypedDict):
48 """A Home Assistant entity that can be used as a player control."""
49
50 entity_id: str
51 # the entity's friendly name, falling back to its entity ID when it has none
52 name: str
53 power: bool
54 volume: bool
55 mute: bool
56
57
58# Selects the entities that can serve each control role.
59CONTROL_TYPE_CAPABILITIES: Final[dict[str, Callable[[HassControlEntity], bool]]] = {
60 CONF_POWER_CONTROLS: lambda entity: entity["power"],
61 CONF_VOLUME_CONTROLS: lambda entity: entity["volume"],
62 CONF_MUTE_CONTROLS: lambda entity: entity["mute"],
63}
64
65
66class HassControlEntityGroup(TypedDict):
67 """
68 The control entities that share a device and an area.
69
70 A device holding entities that sit in different areas is reported as one group per area,
71 since Home Assistant lets an entity override the area it would inherit from its device.
72 Entities without a device are grouped by area alone.
73 """
74
75 device_id: str | None
76 # None for entities that belong to no device, respectively to no area
77 device_name: str | None
78 area_name: str | None
79 entities: list[HassControlEntity]
80
81
82class HassControlEntitySearchResult(TypedDict):
83 """The outcome of a player control entity search."""
84
85 groups: list[HassControlEntityGroup]
86 # True when matches were left out to honor the requested limit
87 truncated: bool
88
89
90class ControlEntitySearch:
91 """Searchable view on the Home Assistant entities that can act as a player control."""
92
93 def __init__(self, provider: HomeAssistantProvider) -> None:
94 """
95 Initialize the search.
96
97 :param provider: The Home Assistant provider to read the registries and states from;
98 its client must be connected before the first search.
99 """
100 self._provider = provider
101 self._lock = asyncio.Lock()
102 self._entries: dict[tuple[str, ...], _CandidateCacheEntry] = {}
103 self._closed = False
104
105 async def search(
106 self,
107 search: str | None = None,
108 control_type: str | None = None,
109 limit: int = SEARCH_CONTROL_ENTITIES_LIMIT,
110 ) -> HassControlEntitySearchResult:
111 """
112 Search the Home Assistant entities that can be used as a player control.
113
114 Music Assistant's own players are never part of the result. Consecutive searches are
115 served from a short lived cache that an entity registry change drops right away, so a
116 newly added or removed entity shows up immediately, while a device or area rename can
117 lag by up to a minute.
118
119 :param search: Text to match, case insensitively, against the entity ID, the entity
120 name, its device name and its area name. Every whitespace separated word must
121 match one of those fields, though not necessarily the same one. All eligible
122 entities match when omitted.
123 :param control_type: Restrict the result to entities that can serve this control role,
124 given as one of the provider's control config keys (``power_controls``,
125 ``volume_controls`` or ``mute_controls``). All roles are returned when omitted.
126 :param limit: Maximum number of entities (not groups) to return, itself capped at
127 ``SEARCH_CONTROL_ENTITIES_MAX_LIMIT``.
128 :return: The matching entities grouped by the device and area they belong to, ordered
129 by area, device and entity name, plus a flag telling whether matches were left out
130 to honor the limit.
131 :raises InvalidDataError: When the control type is unknown or the limit is below one.
132 :raises ProviderUnavailableError: When the provider is no longer loaded.
133 """
134 self._check_open()
135 if control_type is not None and control_type not in CONTROL_TYPE_DOMAINS:
136 msg = f"Invalid control type: {control_type}"
137 raise InvalidDataError(msg)
138 if limit < 1:
139 msg = f"Invalid limit: {limit}"
140 raise InvalidDataError(msg)
141 limit = min(limit, SEARCH_CONTROL_ENTITIES_MAX_LIMIT)
142 domains = CONTROL_DOMAINS if control_type is None else CONTROL_TYPE_DOMAINS[control_type]
143 matches = await self._get_candidates(domains)
144 if control_type is not None:
145 has_capability = CONTROL_TYPE_CAPABILITIES[control_type]
146 matches = [match for match in matches if has_capability(match.entity)]
147 if tokens := (search or "").casefold().split():
148 matches = [match for match in matches if match.matches(tokens)]
149 groups: dict[tuple[str | None, str | None], HassControlEntityGroup] = {}
150 for match in matches[:limit]:
151 group_key = (match.device_id, match.area_id)
152 if (group := groups.get(group_key)) is None:
153 group = HassControlEntityGroup(
154 device_id=match.device_id,
155 device_name=match.device_name,
156 area_name=match.area_name,
157 entities=[],
158 )
159 groups[group_key] = group
160 # the candidates outlive this response, so hand out a copy a caller may mutate
161 group["entities"].append(match.entity.copy())
162 return HassControlEntitySearchResult(
163 groups=list(groups.values()), truncated=len(matches) > limit
164 )
165
166 def close(self) -> None:
167 """Drop everything cached and refuse any further search."""
168 self._closed = True
169 self._entries.clear()
170
171 async def _get_candidates(self, domains: tuple[str, ...]) -> list[_ControlEntityMatch]:
172 """
173 Return the player control candidates found in the given entity domains.
174
175 :param domains: The entity domains to consider.
176 """
177 if (matches := self._lookup(domains)) is not None:
178 return matches
179 async with self._lock:
180 # a concurrent search may have resolved the same domains while this one waited
181 if (matches := self._lookup(domains)) is not None:
182 return matches
183 generation = self._generation()
184 matches = await self._resolve(domains)
185 # candidates resolved against a registry that changed while the sweep was in
186 # flight are stale on arrival, and a provider that unloaded meanwhile must not
187 # see its cache refilled: serve this caller either way, but do not store
188 if not self._closed and generation == self._generation():
189 self._entries[domains] = _CandidateCacheEntry(
190 expires_at=time.monotonic() + CONTROL_ENTITY_CACHE_TTL,
191 generation=generation,
192 matches=matches,
193 )
194 return matches
195
196 async def _resolve(self, domains: tuple[str, ...]) -> list[_ControlEntityMatch]:
197 """
198 Return the entities of the given domains that can serve as a player control.
199
200 :param domains: The entity domains to consider.
201 :return: The candidates in presentation order, each carrying every control role it
202 can serve plus the device and area it belongs to.
203 """
204 entity_registry = await self._provider.get_entity_registry()
205 devices = await self._provider.get_device_registry()
206 areas = await self._provider.get_area_registry()
207 states = await self._provider.get_states(domains=domains)
208 self._check_open()
209 matches: list[_ControlEntityMatch] = []
210 for state in states:
211 capabilities = get_control_capabilities(state, self._provider.logger)
212 if not any(capabilities):
213 continue
214 entity_id = state["entity_id"]
215 registry_entry = entity_registry.get(entity_id)
216 device_id = registry_entry.device_id if registry_entry else None
217 device = devices.get(device_id) if device_id else None
218 # Home Assistant lets an entity override the area it inherits from its device
219 area_id = (registry_entry.area_id if registry_entry else None) or (
220 device["area_id"] if device else None
221 )
222 area = areas.get(area_id or "")
223 name = state["attributes"].get("friendly_name") or entity_id
224 device_name = (device["name_by_user"] or device["name"]) if device else None
225 area_name = area["name"] if area else None
226 matches.append(
227 _ControlEntityMatch(
228 device_id=device_id,
229 device_name=device_name,
230 area_id=area_id,
231 area_name=area_name,
232 entity=HassControlEntity(
233 entity_id=entity_id,
234 name=name,
235 power=capabilities.power,
236 volume=capabilities.volume,
237 mute=capabilities.mute,
238 ),
239 search_text="\n".join(
240 field.casefold()
241 for field in (entity_id, name, device_name, area_name)
242 if field
243 ),
244 )
245 )
246 matches.sort(key=lambda match: match.sort_key)
247 return matches
248
249 def _lookup(self, domains: tuple[str, ...]) -> list[_ControlEntityMatch] | None:
250 """Return the cached candidates of the given domains, None when there are none left."""
251 if (entry := self._entries.get(domains)) is None:
252 return None
253 if entry.generation != self._generation() or entry.expires_at <= time.monotonic():
254 del self._entries[domains]
255 return None
256 return entry.matches
257
258 def _generation(self) -> int:
259 """Return the generation of the entity registry the candidates are resolved against."""
260 # the device and area registries carry no such marker, and the provider already
261 # holds them behind a window of their own that outlives the candidates, so the
262 # entity registry is the only input a fresher sweep could actually improve on
263 return self._provider.entity_registry_generation
264
265 def _check_open(self) -> None:
266 """Raise when the provider this search belongs to is no longer loaded."""
267 if self._closed:
268 msg = "The Home Assistant provider is no longer loaded"
269 raise ProviderUnavailableError(msg)
270
271
272class _ControlEntityMatch(NamedTuple):
273 """A control entity together with the device and area it is presented under."""
274
275 device_id: str | None
276 device_name: str | None
277 area_id: str | None
278 area_name: str | None
279 entity: HassControlEntity
280 # all searchable fields, case folded and joined on a newline. a search token can never
281 # contain whitespace, so a token found in here always sits within a single field
282 search_text: str
283
284 @property
285 def sort_key(self) -> tuple[bool, str, bool, str, str, str]:
286 """Return the ranking key, placing entities without an area or device last."""
287 return (
288 self.area_name is None,
289 (self.area_name or "").casefold(),
290 self.device_name is None,
291 (self.device_name or "").casefold(),
292 self.entity["name"].casefold(),
293 self.entity["entity_id"],
294 )
295
296 def matches(self, tokens: list[str]) -> bool:
297 """
298 Return whether every search token occurs in one of the searchable fields.
299
300 :param tokens: The case folded search tokens.
301 """
302 return all(token in self.search_text for token in tokens)
303
304
305class _CandidateCacheEntry(NamedTuple):
306 """Cached control entity candidates together with what makes them go stale."""
307
308 expires_at: float
309 generation: int
310 matches: list[_ControlEntityMatch]
311