-
-
Notifications
You must be signed in to change notification settings - Fork 28.8k
/
entity.py
343 lines (293 loc) · 11.3 KB
/
entity.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
"""Support for esphome entities."""
from __future__ import annotations
from collections.abc import Awaitable, Callable, Coroutine
import functools
import math
from typing import TYPE_CHECKING, Any, Concatenate, Generic, ParamSpec, TypeVar, cast
from aioesphomeapi import (
APIConnectionError,
EntityCategory as EsphomeEntityCategory,
EntityInfo,
EntityState,
build_unique_id,
)
import voluptuous as vol
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers import entity_platform
import homeassistant.helpers.config_validation as cv
import homeassistant.helpers.device_registry as dr
from homeassistant.helpers.device_registry import DeviceInfo
from homeassistant.helpers.entity import Entity
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from .domain_data import DomainData
# Import config flow so that it's added to the registry
from .entry_data import RuntimeEntryData
from .enum_mapper import EsphomeEnumMapper
_R = TypeVar("_R")
_P = ParamSpec("_P")
_InfoT = TypeVar("_InfoT", bound=EntityInfo)
_EntityT = TypeVar("_EntityT", bound="EsphomeEntity[Any,Any]")
_StateT = TypeVar("_StateT", bound=EntityState)
@callback
def async_static_info_updated(
hass: HomeAssistant,
entry_data: RuntimeEntryData,
platform: entity_platform.EntityPlatform,
async_add_entities: AddEntitiesCallback,
info_type: type[_InfoT],
entity_type: type[_EntityT],
state_type: type[_StateT],
infos: list[EntityInfo],
) -> None:
"""Update entities of this platform when entities are listed."""
current_infos = entry_data.info[info_type]
new_infos: dict[int, EntityInfo] = {}
add_entities: list[_EntityT] = []
for info in infos:
if not current_infos.pop(info.key, None):
# Create new entity
entity = entity_type(entry_data, platform.domain, info, state_type)
add_entities.append(entity)
new_infos[info.key] = info
# Anything still in current_infos is now gone
if current_infos:
device_info = entry_data.device_info
if TYPE_CHECKING:
assert device_info is not None
entry_data.async_remove_entities(
hass, current_infos.values(), device_info.mac_address
)
# Then update the actual info
entry_data.info[info_type] = new_infos
if new_infos:
entry_data.async_update_entity_infos(new_infos.values())
if add_entities:
# Add entities to Home Assistant
async_add_entities(add_entities)
async def platform_async_setup_entry(
hass: HomeAssistant,
entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
*,
info_type: type[_InfoT],
entity_type: type[_EntityT],
state_type: type[_StateT],
) -> None:
"""Set up an esphome platform.
This method is in charge of receiving, distributing and storing
info and state updates.
"""
entry_data: RuntimeEntryData = DomainData.get(hass).get_entry_data(entry)
entry_data.info[info_type] = {}
entry_data.state.setdefault(state_type, {})
platform = entity_platform.async_get_current_platform()
on_static_info_update = functools.partial(
async_static_info_updated,
hass,
entry_data,
platform,
async_add_entities,
info_type,
entity_type,
state_type,
)
entry_data.cleanup_callbacks.append(
entry_data.async_register_static_info_callback(
info_type,
on_static_info_update,
)
)
def esphome_state_property(
func: Callable[[_EntityT], _R],
) -> Callable[[_EntityT], _R | None]:
"""Wrap a state property of an esphome entity.
This checks if the state object in the entity is set, and
prevents writing NAN values to the Home Assistant state machine.
"""
@functools.wraps(func)
def _wrapper(self: _EntityT) -> _R | None:
if not self._has_state:
return None
val = func(self)
if isinstance(val, float) and not math.isfinite(val):
# Home Assistant doesn't use NaN or inf values in state machine
# (not JSON serializable)
return None
return val
return _wrapper
def convert_api_error_ha_error(
func: Callable[Concatenate[_EntityT, _P], Awaitable[None]],
) -> Callable[Concatenate[_EntityT, _P], Coroutine[Any, Any, None]]:
"""Decorate ESPHome command calls that send commands/make changes to the device.
A decorator that wraps the passed in function, catches APIConnectionError errors,
and raises a HomeAssistant error instead.
"""
async def handler(self: _EntityT, *args: _P.args, **kwargs: _P.kwargs) -> None:
try:
return await func(self, *args, **kwargs)
except APIConnectionError as error:
raise HomeAssistantError(
f"Error communicating with device: {error}"
) from error
return handler
ICON_SCHEMA = vol.Schema(cv.icon)
ENTITY_CATEGORIES: EsphomeEnumMapper[EsphomeEntityCategory, EntityCategory | None] = (
EsphomeEnumMapper(
{
EsphomeEntityCategory.NONE: None,
EsphomeEntityCategory.CONFIG: EntityCategory.CONFIG,
EsphomeEntityCategory.DIAGNOSTIC: EntityCategory.DIAGNOSTIC,
}
)
)
class EsphomeEntity(Entity, Generic[_InfoT, _StateT]):
"""Define a base esphome entity."""
_attr_should_poll = False
_static_info: _InfoT
_state: _StateT
_has_state: bool
def __init__(
self,
entry_data: RuntimeEntryData,
domain: str,
entity_info: EntityInfo,
state_type: type[_StateT],
) -> None:
"""Initialize."""
self._entry_data = entry_data
self._on_entry_data_changed()
self._key = entity_info.key
self._state_type = state_type
self._on_static_info_update(entity_info)
assert entry_data.device_info is not None
device_info = entry_data.device_info
self._device_info = device_info
self._attr_device_info = DeviceInfo(
connections={(dr.CONNECTION_NETWORK_MAC, device_info.mac_address)}
)
#
# If `friendly_name` is set, we use the Friendly naming rules, if
# `friendly_name` is not set we make an exception to the naming rules for
# backwards compatibility and use the Legacy naming rules.
#
# Friendly naming
# - Friendly name is prepended to entity names
# - Device Name is prepended to entity ids
# - Entity id is constructed from device name and object id
#
# Legacy naming
# - Device name is not prepended to entity names
# - Device name is not prepended to entity ids
# - Entity id is constructed from entity name
#
if not device_info.friendly_name:
return
self._attr_has_entity_name = True
self.entity_id = f"{domain}.{device_info.name}_{entity_info.object_id}"
async def async_added_to_hass(self) -> None:
"""Register callbacks."""
entry_data = self._entry_data
self.async_on_remove(
entry_data.async_subscribe_device_updated(
self._on_device_update,
)
)
self.async_on_remove(
entry_data.async_subscribe_state_update(
self._state_type, self._key, self._on_state_update
)
)
self.async_on_remove(
entry_data.async_register_key_static_info_updated_callback(
self._static_info, self._on_static_info_update
)
)
self._update_state_from_entry_data()
@callback
def _on_static_info_update(self, static_info: EntityInfo) -> None:
"""Save the static info for this entity when it changes.
This method can be overridden in child classes to know
when the static info changes.
"""
device_info = self._entry_data.device_info
if TYPE_CHECKING:
static_info = cast(_InfoT, static_info)
assert device_info
self._static_info = static_info
self._attr_unique_id = build_unique_id(device_info.mac_address, static_info)
self._attr_entity_registry_enabled_default = not static_info.disabled_by_default
self._attr_name = static_info.name
if entity_category := static_info.entity_category:
self._attr_entity_category = ENTITY_CATEGORIES.from_esphome(entity_category)
else:
self._attr_entity_category = None
if icon := static_info.icon:
self._attr_icon = cast(str, ICON_SCHEMA(icon))
else:
self._attr_icon = None
@callback
def _update_state_from_entry_data(self) -> None:
"""Update state from entry data."""
state = self._entry_data.state
key = self._key
state_type = self._state_type
if has_state := key in state[state_type]:
self._state = cast(_StateT, state[state_type][key])
self._has_state = has_state
@callback
def _on_state_update(self) -> None:
"""Call when state changed.
Behavior can be changed in child classes
"""
self._update_state_from_entry_data()
self.async_write_ha_state()
@callback
def _on_entry_data_changed(self) -> None:
entry_data = self._entry_data
self._api_version = entry_data.api_version
self._client = entry_data.client
@callback
def _on_device_update(self) -> None:
"""Call when device updates or entry data changes."""
self._on_entry_data_changed()
if not self._entry_data.available:
# Only write state if the device has gone unavailable
# since _on_state_update will be called if the device
# is available when the full state arrives
# through the next entity state packet.
self.async_write_ha_state()
@property
def available(self) -> bool:
"""Return if the entity is available."""
if self._device_info.has_deep_sleep:
# During deep sleep the ESP will not be connectable (by design)
# For these cases, show it as available
return self._entry_data.expected_disconnect
return self._entry_data.available
class EsphomeAssistEntity(Entity):
"""Define a base entity for Assist Pipeline entities."""
_attr_has_entity_name = True
_attr_should_poll = False
def __init__(self, entry_data: RuntimeEntryData) -> None:
"""Initialize the binary sensor."""
self._entry_data: RuntimeEntryData = entry_data
assert entry_data.device_info is not None
device_info = entry_data.device_info
self._device_info = device_info
self._attr_unique_id = (
f"{device_info.mac_address}-{self.entity_description.key}"
)
self._attr_device_info = DeviceInfo(
connections={(dr.CONNECTION_NETWORK_MAC, device_info.mac_address)}
)
async def async_added_to_hass(self) -> None:
"""Register update callback."""
await super().async_added_to_hass()
self.async_on_remove(
self._entry_data.async_subscribe_assist_pipeline_update(
self.async_write_ha_state
)
)