From ba284c0d27ece618ca75888fb47790c8a786518e Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Fri, 7 May 2021 23:04:00 +0200
Subject: [PATCH] Add sensor state_class property (#50063)

* Add sensor state_class property

* STATE_CLASS_LATEST -> STATE_CLASS_MEASUREMENT

* Export sensor.state_class in capability_attributes

* Add STATE_CLASS_UNKNOWN

* Fix typing

* Update tests

* STATE_CLASS_UNKNOWN -> STATE_CLASS_OTHER

* Update homeassistant/components/sensor/__init__.py

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>

* Remove STATE_CLASS_OTHER

* Update tests

* Revert test changes

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>
---
 homeassistant/components/hue/sensor_base.py |  6 +++
 homeassistant/components/sensor/__init__.py | 42 +++++++++++++++++----
 2 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/homeassistant/components/hue/sensor_base.py b/homeassistant/components/hue/sensor_base.py
index 9f764e04d28..b8e2af138b2 100644
--- a/homeassistant/components/hue/sensor_base.py
+++ b/homeassistant/components/hue/sensor_base.py
@@ -6,6 +6,7 @@ from aiohue import AiohueException, Unauthorized
 from aiohue.sensors import TYPE_ZLL_PRESENCE
 import async_timeout
 
+from homeassistant.components.sensor import STATE_CLASS_MEASUREMENT
 from homeassistant.core import callback
 from homeassistant.helpers import debounce, entity
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
@@ -177,6 +178,11 @@ class GenericHueSensor(GenericHueDevice, entity.Entity):
             or self.sensor.config.get("reachable", True)
         )
 
+    @property
+    def state_class(self):
+        """Return the state class of this entity, from STATE_CLASSES, if any."""
+        return STATE_CLASS_MEASUREMENT
+
     async def async_added_to_hass(self):
         """When entity is added to hass."""
         self.async_on_remove(
diff --git a/homeassistant/components/sensor/__init__.py b/homeassistant/components/sensor/__init__.py
index 0012b1a3aa2..ed88ca55ceb 100644
--- a/homeassistant/components/sensor/__init__.py
+++ b/homeassistant/components/sensor/__init__.py
@@ -1,10 +1,14 @@
 """Component to interface with various sensors that can be monitored."""
+from __future__ import annotations
 
+from collections.abc import Mapping
 from datetime import timedelta
 import logging
+from typing import Any, cast
 
 import voluptuous as vol
 
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     DEVICE_CLASS_BATTERY,
     DEVICE_CLASS_CO,
@@ -21,17 +25,19 @@ from homeassistant.const import (
     DEVICE_CLASS_TIMESTAMP,
     DEVICE_CLASS_VOLTAGE,
 )
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers.config_validation import (  # noqa: F401
     PLATFORM_SCHEMA,
     PLATFORM_SCHEMA_BASE,
 )
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.entity_component import EntityComponent
-
-# mypy: allow-untyped-defs, no-check-untyped-defs
+from homeassistant.helpers.typing import ConfigType
 
 _LOGGER = logging.getLogger(__name__)
 
+ATTR_STATE_CLASS = "state_class"
+
 DOMAIN = "sensor"
 
 ENTITY_ID_FORMAT = DOMAIN + ".{}"
@@ -56,8 +62,15 @@ DEVICE_CLASSES = [
 
 DEVICE_CLASSES_SCHEMA = vol.All(vol.Lower, vol.In(DEVICE_CLASSES))
 
+# The state represents a measurement in present time
+STATE_CLASS_MEASUREMENT = "measurement"
 
-async def async_setup(hass, config):
+STATE_CLASSES = [STATE_CLASS_MEASUREMENT]
+
+STATE_CLASSES_SCHEMA = vol.All(vol.Lower, vol.In(STATE_CLASSES))
+
+
+async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Track states and offer events for sensors."""
     component = hass.data[DOMAIN] = EntityComponent(
         _LOGGER, DOMAIN, hass, SCAN_INTERVAL
@@ -67,15 +80,30 @@ async def async_setup(hass, config):
     return True
 
 
-async def async_setup_entry(hass, entry):
+async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up a config entry."""
-    return await hass.data[DOMAIN].async_setup_entry(entry)
+    component = cast(EntityComponent, hass.data[DOMAIN])
+    return await component.async_setup_entry(entry)
 
 
-async def async_unload_entry(hass, entry):
+async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    return await hass.data[DOMAIN].async_unload_entry(entry)
+    component = cast(EntityComponent, hass.data[DOMAIN])
+    return await component.async_unload_entry(entry)
 
 
 class SensorEntity(Entity):
     """Base class for sensor entities."""
+
+    @property
+    def state_class(self) -> str | None:
+        """Return the state class of this entity, from STATE_CLASSES, if any."""
+        return None
+
+    @property
+    def capability_attributes(self) -> Mapping[str, Any] | None:
+        """Return the capability attributes."""
+        if self.state_class:
+            return {ATTR_STATE_CLASS: self.state_class}
+
+        return None