Simplify native value and attributes properties for Flu Near You (#56665)

This commit is contained in:
Aaron Bach 2021-09-26 03:12:54 -06:00 committed by GitHub
parent aeba3a703f
commit 01e03a223b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -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)