Modernize template sensor (home-assistant#157251)

Author: Petro31

homeassistant/components/template/sensor.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from datetime import date, datetime
+from datetime import datetime
 import logging
 from typing import Any
 
@@ -18,7 +18,6 @@
     STATE_CLASSES_SCHEMA,
     RestoreSensor,
     SensorDeviceClass,
-    SensorEntity,
     SensorStateClass,
 )
 from homeassistant.components.sensor.helpers import (  # pylint: disable=hass-component-root-import
@@ -35,8 +34,6 @@
     CONF_NAME,
     CONF_SENSORS,
     CONF_STATE,
-    CONF_TRIGGER,
-    CONF_TRIGGERS,
     CONF_UNIQUE_ID,
     CONF_UNIT_OF_MEASUREMENT,
     CONF_VALUE_TEMPLATE,
@@ -54,6 +51,7 @@
 from homeassistant.util import dt as dt_util
 
 from . import TriggerUpdateCoordinator
+from .entity import AbstractTemplateEntity
 from .helpers import (
     async_setup_template_entry,
     async_setup_template_platform,
@@ -136,33 +134,8 @@ def validate_last_reset(val):
     .extend(TEMPLATE_ENTITY_AVAILABILITY_SCHEMA_LEGACY.schema),
 )
 
-
-def extra_validation_checks(val):
-    """Run extra validation checks."""
-    if CONF_TRIGGERS in val or CONF_TRIGGER in val:
-        raise vol.Invalid(
-            "You can only add triggers to template entities if they are defined under"
-            " `template:`. See the template documentation for more information:"
-            " https://www.home-assistant.io/integrations/template/"
-        )
-
-    if CONF_SENSORS not in val and SENSOR_DOMAIN not in val:
-        raise vol.Invalid(f"Required key {SENSOR_DOMAIN} not defined")
-
-    return val
-
-
-PLATFORM_SCHEMA = vol.All(
-    SENSOR_PLATFORM_SCHEMA.extend(
-        {
-            vol.Optional(CONF_TRIGGER): cv.match_all,  # to raise custom warning
-            vol.Optional(CONF_TRIGGERS): cv.match_all,  # to raise custom warning
-            vol.Required(CONF_SENSORS): cv.schema_with_slug_keys(
-                SENSOR_LEGACY_YAML_SCHEMA
-            ),
-        }
-    ),
-    extra_validation_checks,
+PLATFORM_SCHEMA = SENSOR_PLATFORM_SCHEMA.extend(
+    {vol.Required(CONF_SENSORS): cv.schema_with_slug_keys(SENSOR_LEGACY_YAML_SCHEMA)}
 )
 
 _LOGGER = logging.getLogger(__name__)
@@ -213,7 +186,53 @@ def async_create_preview_sensor(
     )
 
 
-class StateSensorEntity(TemplateEntity, SensorEntity):
+class AbstractTemplateSensor(AbstractTemplateEntity, RestoreSensor):
+    """Representation of a template sensor features."""
+
+    _entity_id_format = ENTITY_ID_FORMAT
+
+    # The super init is not called because TemplateEntity and TriggerEntity will call AbstractTemplateEntity.__init__.
+    # This ensures that the __init__ on AbstractTemplateEntity is not called twice.
+    def __init__(self, config: ConfigType) -> None:  # pylint: disable=super-init-not-called
+        """Initialize the features."""
+        self._attr_native_unit_of_measurement = config.get(CONF_UNIT_OF_MEASUREMENT)
+        self._attr_device_class = config.get(CONF_DEVICE_CLASS)
+        self._attr_state_class = config.get(CONF_STATE_CLASS)
+        self._template: template.Template = config[CONF_STATE]
+        self._attr_last_reset_template: template.Template | None = config.get(
+            ATTR_LAST_RESET
+        )
+
+    @callback
+    def _update_last_reset(self, result: Any) -> None:
+        if isinstance(result, datetime):
+            self._attr_last_reset = result
+            return
+
+        parsed_timestamp = dt_util.parse_datetime(result)
+        if parsed_timestamp is None:
+            _LOGGER.warning(
+                "%s rendered invalid timestamp for last_reset attribute: %s",
+                self.entity_id,
+                result,
+            )
+        else:
+            self._attr_last_reset = parsed_timestamp
+
+    def _handle_state(self, result: Any) -> None:
+        if result is None or self.device_class not in (
+            SensorDeviceClass.DATE,
+            SensorDeviceClass.TIMESTAMP,
+        ):
+            self._attr_native_value = result
+            return
+
+        self._attr_native_value = async_parse_date_datetime(
+            result, self.entity_id, self.device_class
+        )
+
+
+class StateSensorEntity(TemplateEntity, AbstractTemplateSensor):
     """Representation of a Template Sensor."""
 
     _attr_should_poll = False
@@ -226,14 +245,8 @@ def __init__(
         unique_id: str | None,
     ) -> None:
         """Initialize the sensor."""
-        super().__init__(hass, config, unique_id)
-        self._attr_native_unit_of_measurement = config.get(CONF_UNIT_OF_MEASUREMENT)
-        self._attr_device_class = config.get(CONF_DEVICE_CLASS)
-        self._attr_state_class = config.get(CONF_STATE_CLASS)
-        self._template: template.Template = config[CONF_STATE]
-        self._attr_last_reset_template: template.Template | None = config.get(
-            ATTR_LAST_RESET
-        )
+        TemplateEntity.__init__(self, hass, config, unique_id)
+        AbstractTemplateSensor.__init__(self, config)
 
     @callback
     def _async_setup_templates(self) -> None:
@@ -251,35 +264,20 @@ def _async_setup_templates(self) -> None:
 
         super()._async_setup_templates()
 
-    @callback
-    def _update_last_reset(self, result):
-        self._attr_last_reset = result
-
     @callback
     def _update_state(self, result):
         super()._update_state(result)
         if isinstance(result, TemplateError):
             self._attr_native_value = None
             return
 
-        if result is None or self.device_class not in (
-            SensorDeviceClass.DATE,
-            SensorDeviceClass.TIMESTAMP,
-        ):
-            self._attr_native_value = result
-            return
-
-        self._attr_native_value = async_parse_date_datetime(
-            result, self.entity_id, self.device_class
-        )
+        self._handle_state(result)
 
 
-class TriggerSensorEntity(TriggerEntity, RestoreSensor):
+class TriggerSensorEntity(TriggerEntity, AbstractTemplateSensor):
     """Sensor entity based on trigger data."""
 
-    _entity_id_format = ENTITY_ID_FORMAT
     domain = SENSOR_DOMAIN
-    extra_template_keys = (CONF_STATE,)
 
     def __init__(
         self,
@@ -288,18 +286,18 @@ def __init__(
         config: ConfigType,
     ) -> None:
         """Initialize."""
-        super().__init__(hass, coordinator, config)
+        TriggerEntity.__init__(self, hass, coordinator, config)
+        AbstractTemplateSensor.__init__(self, config)
 
+        self._to_render_simple.append(CONF_STATE)
         self._parse_result.add(CONF_STATE)
-        if (last_reset_template := config.get(ATTR_LAST_RESET)) is not None:
+
+        if last_reset_template := self._attr_last_reset_template:
             if last_reset_template.is_static:
                 self._static_rendered[ATTR_LAST_RESET] = last_reset_template.template
             else:
                 self._to_render_simple.append(ATTR_LAST_RESET)
 
-        self._attr_state_class = config.get(CONF_STATE_CLASS)
-        self._attr_native_unit_of_measurement = config.get(CONF_UNIT_OF_MEASUREMENT)
-
     async def async_added_to_hass(self) -> None:
         """Restore last state."""
         await super().async_added_to_hass()
@@ -311,39 +309,17 @@ async def async_added_to_hass(self) -> None:
             # then we should not restore state
             and CONF_STATE not in self._rendered
         ):
-            self._rendered[CONF_STATE] = extra_data.native_value
+            self._attr_native_value = extra_data.native_value
             self.restore_attributes(last_state)
 
-    @property
-    def native_value(self) -> str | datetime | date | None:
-        """Return state of the sensor."""
-        return self._rendered.get(CONF_STATE)
-
     @callback
     def _process_data(self) -> None:
         """Process new data."""
         super()._process_data()
 
         # Update last_reset
-        if ATTR_LAST_RESET in self._rendered:
-            parsed_timestamp = dt_util.parse_datetime(self._rendered[ATTR_LAST_RESET])
-            if parsed_timestamp is None:
-                _LOGGER.warning(
-                    "%s rendered invalid timestamp for last_reset attribute: %s",
-                    self.entity_id,
-                    self._rendered.get(ATTR_LAST_RESET),
-                )
-            else:
-                self._attr_last_reset = parsed_timestamp
+        if (last_reset := self._rendered.get(ATTR_LAST_RESET)) is not None:
+            self._update_last_reset(last_reset)
 
-        if (
-            state := self._rendered.get(CONF_STATE)
-        ) is None or self.device_class not in (
-            SensorDeviceClass.DATE,
-            SensorDeviceClass.TIMESTAMP,
-        ):
-            return
-
-        self._rendered[CONF_STATE] = async_parse_date_datetime(
-            state, self.entity_id, self.device_class
-        )
+        rendered = self._rendered.get(CONF_STATE)
+        self._handle_state(rendered)

tests/components/template/test_sensor.py

@@ -1456,7 +1456,7 @@ async def test_trigger_not_allowed_platform_config(
     state = hass.states.get(TEST_NAME)
     assert state is None
     assert (
-        "You can only add triggers to template entities if they are defined under `template:`."
+        "Invalid config for 'sensor' from integration 'template': 'trigger' is an invalid option for"
         in caplog_setup_text
     )