New template sensor attributes (#26127)

* updated sensor and test files

* Formatting fixes

* Updated attribute template code

* Black formatting

* Code improvements based on feedback on binary_sensor pull request

* Updated tests

* Remove duplicated code and fix tests

* Black formatting on tests

* Remove link from docstring

* Moved default to schema

* Formatting fix and change to use dict[key] to retrieve attribute_templates

homeassistant/components/template/sensor.py

@@ -1,6 +1,7 @@
 """Allows the creation of a sensor that breaks out state_attributes."""
 import logging
 from typing import Optional
+from itertools import chain
 
 import voluptuous as vol
 
@@ -28,6 +29,8 @@ import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity import Entity, async_generate_entity_id
 from homeassistant.helpers.event import async_track_state_change
 
+CONF_ATTRIBUTE_TEMPLATES = "attribute_templates"
+
 _LOGGER = logging.getLogger(__name__)
 
 SENSOR_SCHEMA = vol.Schema(
@@ -36,6 +39,9 @@ SENSOR_SCHEMA = vol.Schema(
         vol.Optional(CONF_ICON_TEMPLATE): cv.template,
         vol.Optional(CONF_ENTITY_PICTURE_TEMPLATE): cv.template,
         vol.Optional(CONF_FRIENDLY_NAME_TEMPLATE): cv.template,
+        vol.Optional(CONF_ATTRIBUTE_TEMPLATES, default={}): vol.Schema(
+            {cv.string: cv.template}
+        ),
         vol.Optional(ATTR_FRIENDLY_NAME): cv.string,
         vol.Optional(ATTR_UNIT_OF_MEASUREMENT): cv.string,
         vol.Optional(CONF_DEVICE_CLASS): DEVICE_CLASSES_SCHEMA,
@@ -60,17 +66,20 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
         friendly_name_template = device_config.get(CONF_FRIENDLY_NAME_TEMPLATE)
         unit_of_measurement = device_config.get(ATTR_UNIT_OF_MEASUREMENT)
         device_class = device_config.get(CONF_DEVICE_CLASS)
+        attribute_templates = device_config[CONF_ATTRIBUTE_TEMPLATES]
 
         entity_ids = set()
         manual_entity_ids = device_config.get(ATTR_ENTITY_ID)
         invalid_templates = []
 
-        for tpl_name, template in (
-            (CONF_VALUE_TEMPLATE, state_template),
-            (CONF_ICON_TEMPLATE, icon_template),
-            (CONF_ENTITY_PICTURE_TEMPLATE, entity_picture_template),
-            (CONF_FRIENDLY_NAME_TEMPLATE, friendly_name_template),
-        ):
+        templates = {
+            CONF_VALUE_TEMPLATE: state_template,
+            CONF_ICON_TEMPLATE: icon_template,
+            CONF_ENTITY_PICTURE_TEMPLATE: entity_picture_template,
+            CONF_FRIENDLY_NAME_TEMPLATE: friendly_name_template,
+        }
+
+        for tpl_name, template in chain(templates.items(), attribute_templates.items()):
             if template is None:
                 continue
             template.hass = hass
@@ -82,7 +91,7 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
             if template_entity_ids == MATCH_ALL:
                 entity_ids = MATCH_ALL
                 # Cut off _template from name
-                invalid_templates.append(tpl_name[:-9])
+                invalid_templates.append(tpl_name.replace("_template", ""))
             elif entity_ids != MATCH_ALL:
                 entity_ids |= set(template_entity_ids)
 
@@ -113,6 +122,7 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
                 entity_picture_template,
                 entity_ids,
                 device_class,
+                attribute_templates,
             )
         )
     if not sensors:
@@ -138,6 +148,7 @@ class SensorTemplate(Entity):
         entity_picture_template,
         entity_ids,
         device_class,
+        attribute_templates,
     ):
         """Initialize the sensor."""
         self.hass = hass
@@ -155,6 +166,8 @@ class SensorTemplate(Entity):
         self._entity_picture = None
         self._entities = entity_ids
         self._device_class = device_class
+        self._attribute_templates = attribute_templates
+        self._attributes = {}
 
     async def async_added_to_hass(self):
         """Register callbacks."""
@@ -209,6 +222,11 @@ class SensorTemplate(Entity):
         """Return the unit_of_measurement of the device."""
         return self._unit_of_measurement
 
+    @property
+    def device_state_attributes(self):
+        """Return the state attributes."""
+        return self._attributes
+
     @property
     def should_poll(self):
         """No polling needed."""
@@ -229,11 +247,23 @@ class SensorTemplate(Entity):
             else:
                 self._state = None
                 _LOGGER.error("Could not render template %s: %s", self._name, ex)
-        for property_name, template in (
-            ("_icon", self._icon_template),
-            ("_entity_picture", self._entity_picture_template),
-            ("_name", self._friendly_name_template),
-        ):
+
+        templates = {
+            "_icon": self._icon_template,
+            "_entity_picture": self._entity_picture_template,
+            "_name": self._friendly_name_template,
+        }
+
+        attrs = {}
+        for key, value in self._attribute_templates.items():
+            try:
+                attrs[key] = value.async_render()
+            except TemplateError as err:
+                _LOGGER.error("Error rendering attribute %s: %s", key, err)
+
+        self._attributes = attrs
+
+        for property_name, template in templates.items():
             if template is None:
                 continue