Allow defining state class for template sensors (#52130)

This commit is contained in:
Paulus Schoutsen 2021-06-23 14:37:04 -07:00 committed by GitHub
parent 6352d8fb0e
commit cc00617cd5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 70 additions and 73 deletions

View file

@ -112,6 +112,8 @@ class _TemplateAttribute:
class TemplateEntity(Entity):
"""Entity that uses templates to calculate attributes."""
_attr_should_poll = False
def __init__(
self,
*,
@ -124,54 +126,27 @@ class TemplateEntity(Entity):
self._template_attrs = {}
self._async_update = None
self._attribute_templates = attribute_templates
self._attributes = {}
self._attr_extra_state_attributes = {}
self._availability_template = availability_template
self._available = True
self._attr_available = True
self._icon_template = icon_template
self._entity_picture_template = entity_picture_template
self._icon = None
self._entity_picture = None
self._self_ref_update_count = 0
@property
def should_poll(self):
"""No polling needed."""
return False
@callback
def _update_available(self, result):
if isinstance(result, TemplateError):
self._available = True
self._attr_available = True
return
self._available = result_as_boolean(result)
self._attr_available = result_as_boolean(result)
@callback
def _update_state(self, result):
if self._availability_template:
return
self._available = not isinstance(result, TemplateError)
@property
def available(self) -> bool:
"""Return if the device is available."""
return self._available
@property
def icon(self):
"""Return the icon to use in the frontend, if any."""
return self._icon
@property
def entity_picture(self):
"""Return the entity_picture to use in the frontend, if any."""
return self._entity_picture
@property
def extra_state_attributes(self):
"""Return the state attributes."""
return self._attributes
self._attr_available = not isinstance(result, TemplateError)
@callback
def _add_attribute_template(self, attribute_key, attribute_template):
@ -179,7 +154,7 @@ class TemplateEntity(Entity):
def _update_attribute(result):
attr_result = None if isinstance(result, TemplateError) else result
self._attributes[attribute_key] = attr_result
self._attr_extra_state_attributes[attribute_key] = attr_result
self.add_template_attribute(
attribute_key, attribute_template, None, _update_attribute
@ -271,18 +246,21 @@ class TemplateEntity(Entity):
"""Run when entity about to be added to hass."""
if self._availability_template is not None:
self.add_template_attribute(
"_available", self._availability_template, None, self._update_available
"_attr_available",
self._availability_template,
None,
self._update_available,
)
if self._attribute_templates is not None:
for key, value in self._attribute_templates.items():
self._add_attribute_template(key, value)
if self._icon_template is not None:
self.add_template_attribute(
"_icon", self._icon_template, vol.Or(cv.whitespace, cv.icon)
"_attr_icon", self._icon_template, vol.Or(cv.whitespace, cv.icon)
)
if self._entity_picture_template is not None:
self.add_template_attribute(
"_entity_picture", self._entity_picture_template
"_attr_entity_picture", self._entity_picture_template
)
if self.hass.state == CoreState.running:
await self._async_template_startup()