Move attributes to be excluded from recording to entity classes (#100239)

Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/components/automation/__init__.py

@@ -314,6 +314,9 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
 class BaseAutomationEntity(ToggleEntity, ABC):
     """Base class for automation entities."""
 
+    _entity_component_unrecorded_attributes = frozenset(
+        (ATTR_LAST_TRIGGERED, ATTR_MODE, ATTR_CUR, ATTR_MAX, CONF_ID)
+    )
     raw_config: ConfigType | None
 
     @property

homeassistant/components/automation/recorder.py

@@ -1,12 +0,0 @@
-"""Integration platform for recorder."""
-from __future__ import annotations
-
-from homeassistant.core import HomeAssistant, callback
-
-from . import ATTR_CUR, ATTR_LAST_TRIGGERED, ATTR_MAX, ATTR_MODE, CONF_ID
-
-
-@callback
-def exclude_attributes(hass: HomeAssistant) -> set[str]:
-    """Exclude extra attributes from being recorded in the database."""
-    return {ATTR_LAST_TRIGGERED, ATTR_MODE, ATTR_CUR, ATTR_MAX, CONF_ID}

homeassistant/components/recorder/db_schema.py

@@ -576,6 +576,8 @@ class StateAttributes(Base):
             integration_attrs := exclude_attrs_by_domain.get(entity_info["domain"])
         ):
             exclude_attrs |= integration_attrs
+        if state_info := state.state_info:
+            exclude_attrs |= state_info["unrecorded_attributes"]
         encoder = json_bytes_strip_null if dialect == PSQL_DIALECT else json_bytes
         bytes_result = encoder(
             {k: v for k, v in state.attributes.items() if k not in exclude_attrs}

homeassistant/core.py

@@ -95,6 +95,7 @@ if TYPE_CHECKING:
     from .auth import AuthManager
     from .components.http import ApiConfig, HomeAssistantHTTP
     from .config_entries import ConfigEntries
+    from .helpers.entity import StateInfo
 
 
 STAGE_1_SHUTDOWN_TIMEOUT = 100
@@ -1249,6 +1250,7 @@ class State:
         last_updated: datetime.datetime | None = None,
         context: Context | None = None,
         validate_entity_id: bool | None = True,
+        state_info: StateInfo | None = None,
     ) -> None:
         """Initialize a new state."""
         state = str(state)
@@ -1267,6 +1269,7 @@ class State:
         self.last_updated = last_updated or dt_util.utcnow()
         self.last_changed = last_changed or self.last_updated
         self.context = context or Context()
+        self.state_info = state_info
         self.domain, self.object_id = split_entity_id(self.entity_id)
         self._as_dict: ReadOnlyDict[str, Collection[Any]] | None = None
 
@@ -1637,6 +1640,7 @@ class StateMachine:
         attributes: Mapping[str, Any] | None = None,
         force_update: bool = False,
         context: Context | None = None,
+        state_info: StateInfo | None = None,
     ) -> None:
         """Set the state of an entity, add entity if it does not exist.
 
@@ -1688,6 +1692,7 @@ class StateMachine:
             now,
             context,
             old_state is None,
+            state_info,
         )
         if old_state is not None:
             old_state.expire()

homeassistant/helpers/entity.py

@@ -201,6 +201,12 @@ class EntityInfo(TypedDict):
     config_entry: NotRequired[str]
 
 
+class StateInfo(TypedDict):
+    """State info."""
+
+    unrecorded_attributes: frozenset[str]
+
+
 class EntityPlatformState(Enum):
     """The platform state of an entity."""
 
@@ -297,6 +303,22 @@ class Entity(ABC):
     # If entity is added to an entity platform
     _platform_state = EntityPlatformState.NOT_ADDED
 
+    # Attributes to exclude from recording, only set by base components, e.g. light
+    _entity_component_unrecorded_attributes: frozenset[str] = frozenset()
+    # Additional integration specific attributes to exclude from recording, set by
+    # platforms, e.g. a derived class in hue.light
+    _unrecorded_attributes: frozenset[str] = frozenset()
+    # Union of _entity_component_unrecorded_attributes and _unrecorded_attributes,
+    # set automatically by __init_subclass__
+    __combined_unrecorded_attributes: frozenset[str] = (
+        _entity_component_unrecorded_attributes | _unrecorded_attributes
+    )
+
+    # StateInfo. Set by EntityPlatform by calling async_internal_added_to_hass
+    # While not purely typed, it makes typehinting more useful for us
+    # and removes the need for constant None checks or asserts.
+    _state_info: StateInfo = None  # type: ignore[assignment]
+
     # Entity Properties
     _attr_assumed_state: bool = False
     _attr_attribution: str | None = None
@@ -321,6 +343,13 @@ class Entity(ABC):
     _attr_unique_id: str | None = None
     _attr_unit_of_measurement: str | None
 
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        """Initialize an Entity subclass."""
+        super().__init_subclass__(**kwargs)
+        cls.__combined_unrecorded_attributes = (
+            cls._entity_component_unrecorded_attributes | cls._unrecorded_attributes
+        )
+
     @property
     def should_poll(self) -> bool:
         """Return True if entity has to be polled for state.
@@ -875,7 +904,12 @@ class Entity(ABC):
 
         try:
             hass.states.async_set(
-                entity_id, state, attr, self.force_update, self._context
+                entity_id,
+                state,
+                attr,
+                self.force_update,
+                self._context,
+                self._state_info,
             )
         except InvalidStateError:
             _LOGGER.exception("Failed to set state, fall back to %s", STATE_UNKNOWN)
@@ -1081,15 +1115,19 @@ class Entity(ABC):
 
         Not to be extended by integrations.
         """
-        info: EntityInfo = {
+        entity_info: EntityInfo = {
             "domain": self.platform.platform_name,
             "custom_component": "custom_components" in type(self).__module__,
         }
 
         if self.platform.config_entry:
-            info["config_entry"] = self.platform.config_entry.entry_id
+            entity_info["config_entry"] = self.platform.config_entry.entry_id
 
-        entity_sources(self.hass)[self.entity_id] = info
+        entity_sources(self.hass)[self.entity_id] = entity_info
+
+        self._state_info = {
+            "unrecorded_attributes": self.__combined_unrecorded_attributes
+        }
 
         if self.registry_entry is not None:
             # This is an assert as it should never happen, but helps in tests