diff --git a/homeassistant/components/flunearyou/sensor.py b/homeassistant/components/flunearyou/sensor.py index 72c0a7b1118..1a7aba5966b 100644 --- a/homeassistant/components/flunearyou/sensor.py +++ b/homeassistant/components/flunearyou/sensor.py @@ -1,6 +1,9 @@ """Support for user- and CDC-based flu info sensors from Flu Near You.""" from __future__ import annotations +from collections.abc import Mapping +from typing import Any, Union, cast + from homeassistant.components.sensor import ( STATE_CLASS_MEASUREMENT, SensorEntity, @@ -13,8 +16,9 @@ from homeassistant.const import ( CONF_LATITUDE, CONF_LONGITUDE, ) -from homeassistant.core import HomeAssistant, callback +from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback +from homeassistant.helpers.typing import StateType from homeassistant.helpers.update_coordinator import ( CoordinatorEntity, DataUpdateCoordinator, @@ -136,6 +140,8 @@ async def async_setup_entry( class FluNearYouSensor(CoordinatorEntity, SensorEntity): """Define a base Flu Near You sensor.""" + DEFAULT_EXTRA_STATE_ATTRIBUTES = {ATTR_ATTRIBUTION: DEFAULT_ATTRIBUTION} + def __init__( self, coordinator: DataUpdateCoordinator, @@ -145,7 +151,6 @@ class FluNearYouSensor(CoordinatorEntity, SensorEntity): """Initialize the sensor.""" super().__init__(coordinator) - self._attr_extra_state_attributes = {ATTR_ATTRIBUTION: DEFAULT_ATTRIBUTION} self._attr_unique_id = ( f"{entry.data[CONF_LATITUDE]}," f"{entry.data[CONF_LONGITUDE]}_{description.key}" @@ -153,68 +158,61 @@ class FluNearYouSensor(CoordinatorEntity, SensorEntity): self._entry = entry self.entity_description = description - @callback - def _handle_coordinator_update(self) -> None: - """Handle updated data from the coordinator.""" - self.update_from_latest_data() - self.async_write_ha_state() - - async def async_added_to_hass(self) -> None: - """Register callbacks.""" - await super().async_added_to_hass() - self.update_from_latest_data() - - @callback - def update_from_latest_data(self) -> None: - """Update the sensor.""" - raise NotImplementedError - class CdcSensor(FluNearYouSensor): """Define a sensor for CDC reports.""" - @callback - def update_from_latest_data(self) -> None: - """Update the sensor.""" - self._attr_extra_state_attributes.update( - { - ATTR_REPORTED_DATE: self.coordinator.data["week_date"], - ATTR_STATE: self.coordinator.data["name"], - } + @property + def extra_state_attributes(self) -> Mapping[str, Any] | None: + """Return entity specific state attributes.""" + return { + **self.DEFAULT_EXTRA_STATE_ATTRIBUTES, + ATTR_REPORTED_DATE: self.coordinator.data["week_date"], + ATTR_STATE: self.coordinator.data["name"], + } + + @property + def native_value(self) -> StateType: + """Return the value reported by the sensor.""" + return cast( + Union[str, None], self.coordinator.data[self.entity_description.key] ) - self._attr_native_value = self.coordinator.data[self.entity_description.key] class UserSensor(FluNearYouSensor): """Define a sensor for user reports.""" - @callback - def update_from_latest_data(self) -> None: - """Update the sensor.""" - self._attr_extra_state_attributes.update( - { - ATTR_CITY: self.coordinator.data["local"]["city"].split("(")[0], - ATTR_REPORTED_LATITUDE: self.coordinator.data["local"]["latitude"], - ATTR_REPORTED_LONGITUDE: self.coordinator.data["local"]["longitude"], - ATTR_STATE: self.coordinator.data["state"]["name"], - ATTR_ZIP_CODE: self.coordinator.data["local"]["zip"], - } - ) + @property + def extra_state_attributes(self) -> Mapping[str, Any] | None: + """Return entity specific state attributes.""" + attrs = { + **self.DEFAULT_EXTRA_STATE_ATTRIBUTES, + ATTR_CITY: self.coordinator.data["local"]["city"].split("(")[0], + ATTR_REPORTED_LATITUDE: self.coordinator.data["local"]["latitude"], + ATTR_REPORTED_LONGITUDE: self.coordinator.data["local"]["longitude"], + ATTR_STATE: self.coordinator.data["state"]["name"], + ATTR_ZIP_CODE: self.coordinator.data["local"]["zip"], + } if self.entity_description.key in self.coordinator.data["state"]["data"]: states_key = self.entity_description.key elif self.entity_description.key in EXTENDED_SENSOR_TYPE_MAPPING: states_key = EXTENDED_SENSOR_TYPE_MAPPING[self.entity_description.key] - self._attr_extra_state_attributes[ - ATTR_STATE_REPORTS_THIS_WEEK - ] = self.coordinator.data["state"]["data"][states_key] - self._attr_extra_state_attributes[ - ATTR_STATE_REPORTS_LAST_WEEK - ] = self.coordinator.data["state"]["last_week_data"][states_key] + attrs[ATTR_STATE_REPORTS_THIS_WEEK] = self.coordinator.data["state"]["data"][ + states_key + ] + attrs[ATTR_STATE_REPORTS_LAST_WEEK] = self.coordinator.data["state"][ + "last_week_data" + ][states_key] + return attrs + + @property + def native_value(self) -> StateType: + """Return the value reported by the sensor.""" if self.entity_description.key == SENSOR_TYPE_USER_TOTAL: - self._attr_native_value = sum( + value = sum( v for k, v in self.coordinator.data["local"].items() if k @@ -227,6 +225,6 @@ class UserSensor(FluNearYouSensor): ) ) else: - self._attr_native_value = self.coordinator.data["local"][ - self.entity_description.key - ] + value = self.coordinator.data["local"][self.entity_description.key] + + return cast(int, value)