Add sensor state_class property

homeassistant/components/hue/sensor_base.py

@@ -6,6 +6,7 @@
 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 @@ def available(self):
             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(

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 @@
     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_SCHEMA = vol.All(vol.Lower, vol.In(DEVICE_CLASSES))
 
+# The state represents a measurement in present time
+STATE_CLASS_MEASUREMENT = "measurement"
+
+STATE_CLASSES = [STATE_CLASS_MEASUREMENT]
+
+STATE_CLASSES_SCHEMA = vol.All(vol.Lower, vol.In(STATE_CLASSES))
 
-async def async_setup(hass, config):
+
+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

tests/components/broadlink/test_sensors.py

@@ -27,7 +27,7 @@ async def test_a1_sensor_setup(hass):
     assert mock_api.check_sensors_raw.call_count == 1
     device_entry = device_registry.async_get_device({(DOMAIN, mock_entry.unique_id)})
     entries = async_entries_for_device(entity_registry, device_entry.id)
-    sensors = {entry for entry in entries if entry.domain == SENSOR_DOMAIN}
+    sensors = [entry for entry in entries if entry.domain == SENSOR_DOMAIN]
     assert len(sensors) == 5
 
     sensors_and_states = {
@@ -62,7 +62,7 @@ async def test_a1_sensor_update(hass):
 
     device_entry = device_registry.async_get_device({(DOMAIN, mock_entry.unique_id)})
     entries = async_entries_for_device(entity_registry, device_entry.id)
-    sensors = {entry for entry in entries if entry.domain == SENSOR_DOMAIN}
+    sensors = [entry for entry in entries if entry.domain == SENSOR_DOMAIN]
     assert len(sensors) == 5
 
     mock_api.check_sensors_raw.return_value = {
@@ -104,7 +104,7 @@ async def test_rm_pro_sensor_setup(hass):
     assert mock_api.check_sensors.call_count == 1
     device_entry = device_registry.async_get_device({(DOMAIN, mock_entry.unique_id)})
     entries = async_entries_for_device(entity_registry, device_entry.id)
-    sensors = {entry for entry in entries if entry.domain == SENSOR_DOMAIN}
+    sensors = [entry for entry in entries if entry.domain == SENSOR_DOMAIN]
     assert len(sensors) == 1
 
     sensors_and_states = {
@@ -127,7 +127,7 @@ async def test_rm_pro_sensor_update(hass):
 
     device_entry = device_registry.async_get_device({(DOMAIN, mock_entry.unique_id)})
     entries = async_entries_for_device(entity_registry, device_entry.id)
-    sensors = {entry for entry in entries if entry.domain == SENSOR_DOMAIN}
+    sensors = [entry for entry in entries if entry.domain == SENSOR_DOMAIN]
     assert len(sensors) == 1
 
     mock_api.check_sensors.return_value = {"temperature": 25.8}
@@ -157,7 +157,7 @@ async def test_rm_mini3_no_sensor(hass):
     assert mock_api.check_sensors.call_count <= 1
     device_entry = device_registry.async_get_device({(DOMAIN, mock_entry.unique_id)})
     entries = async_entries_for_device(entity_registry, device_entry.id)
-    sensors = {entry for entry in entries if entry.domain == SENSOR_DOMAIN}
+    sensors = [entry for entry in entries if entry.domain == SENSOR_DOMAIN]
     assert len(sensors) == 0
 
 
@@ -175,7 +175,7 @@ async def test_rm4_pro_hts2_sensor_setup(hass):
     assert mock_api.check_sensors.call_count == 1
     device_entry = device_registry.async_get_device({(DOMAIN, mock_entry.unique_id)})
     entries = async_entries_for_device(entity_registry, device_entry.id)
-    sensors = {entry for entry in entries if entry.domain == SENSOR_DOMAIN}
+    sensors = [entry for entry in entries if entry.domain == SENSOR_DOMAIN]
     assert len(sensors) == 2
 
     sensors_and_states = {
@@ -201,7 +201,7 @@ async def test_rm4_pro_hts2_sensor_update(hass):
 
     device_entry = device_registry.async_get_device({(DOMAIN, mock_entry.unique_id)})
     entries = async_entries_for_device(entity_registry, device_entry.id)
-    sensors = {entry for entry in entries if entry.domain == SENSOR_DOMAIN}
+    sensors = [entry for entry in entries if entry.domain == SENSOR_DOMAIN]
     assert len(sensors) == 2
 
     mock_api.check_sensors.return_value = {"temperature": 16.8, "humidity": 34.0}
@@ -234,5 +234,5 @@ async def test_rm4_pro_no_sensor(hass):
     assert mock_api.check_sensors.call_count <= 1
     device_entry = device_registry.async_get_device({(DOMAIN, mock_entry.unique_id)})
     entries = async_entries_for_device(entity_registry, device_entry.id)
-    sensors = {entry for entry in entries if entry.domain == SENSOR_DOMAIN}
+    sensors = [entry for entry in entries if entry.domain == SENSOR_DOMAIN]
     assert len(sensors) == 0

tests/components/rest/test_sensor.py

@@ -410,7 +410,11 @@ async def test_update_with_json_attrs_no_data(hass, caplog):
 
     state = hass.states.get("sensor.foo")
     assert state.state == STATE_UNKNOWN
-    assert state.attributes == {"unit_of_measurement": "MB", "friendly_name": "foo"}
+    assert state.attributes == {
+        "unit_of_measurement": "MB",
+        "friendly_name": "foo",
+        "state_class": None,
+    }
     assert "Empty reply" in caplog.text
 
 
@@ -445,7 +449,11 @@ async def test_update_with_json_attrs_not_dict(hass, caplog):
 
     state = hass.states.get("sensor.foo")
     assert state.state == ""
-    assert state.attributes == {"unit_of_measurement": "MB", "friendly_name": "foo"}
+    assert state.attributes == {
+        "unit_of_measurement": "MB",
+        "friendly_name": "foo",
+        "state_class": None,
+    }
     assert "not a dictionary or list" in caplog.text
 
 
@@ -481,7 +489,11 @@ async def test_update_with_json_attrs_bad_JSON(hass, caplog):
 
     state = hass.states.get("sensor.foo")
     assert state.state == STATE_UNKNOWN
-    assert state.attributes == {"unit_of_measurement": "MB", "friendly_name": "foo"}
+    assert state.attributes == {
+        "unit_of_measurement": "MB",
+        "friendly_name": "foo",
+        "state_class": None,
+    }
     assert "Erroneous JSON" in caplog.text