diff --git a/homeassistant/components/trafikverket_train/sensor.py b/homeassistant/components/trafikverket_train/sensor.py index 47f31e35c63..58cf62dd505 100644 --- a/homeassistant/components/trafikverket_train/sensor.py +++ b/homeassistant/components/trafikverket_train/sensor.py @@ -3,9 +3,7 @@ from __future__ import annotations from collections.abc import Callable from dataclasses import dataclass -from datetime import datetime, time, timedelta - -from pytrafikverket.trafikverket_train import StationInfo +from datetime import datetime from homeassistant.components.sensor import ( SensorDeviceClass, @@ -13,37 +11,23 @@ from homeassistant.components.sensor import ( SensorEntityDescription, ) from homeassistant.config_entries import ConfigEntry -from homeassistant.const import CONF_NAME, CONF_WEEKDAY +from homeassistant.const import CONF_NAME, UnitOfTime from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.device_registry import DeviceEntryType from homeassistant.helpers.entity import DeviceInfo from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.typing import StateType from homeassistant.helpers.update_coordinator import CoordinatorEntity -from homeassistant.util import dt as dt_util -from .const import CONF_TIME, DOMAIN +from .const import ATTRIBUTION, DOMAIN from .coordinator import TrainData, TVDataUpdateCoordinator -ATTR_DEPARTURE_STATE = "departure_state" -ATTR_CANCELED = "canceled" -ATTR_DELAY_TIME = "number_of_minutes_delayed" -ATTR_PLANNED_TIME = "planned_time" -ATTR_ESTIMATED_TIME = "estimated_time" -ATTR_ACTUAL_TIME = "actual_time" -ATTR_OTHER_INFORMATION = "other_information" -ATTR_DEVIATIONS = "deviations" - -ICON = "mdi:train" -SCAN_INTERVAL = timedelta(minutes=5) - @dataclass class TrafikverketRequiredKeysMixin: """Mixin for required keys.""" value_fn: Callable[[TrainData], StateType | datetime] - extra_fn: Callable[[TrainData], dict[str, StateType | datetime]] @dataclass @@ -60,16 +44,64 @@ SENSOR_TYPES: tuple[TrafikverketSensorEntityDescription, ...] = ( icon="mdi:clock", device_class=SensorDeviceClass.TIMESTAMP, value_fn=lambda data: data.departure_time, - extra_fn=lambda data: { - ATTR_DEPARTURE_STATE: data.departure_state, - ATTR_CANCELED: data.cancelled, - ATTR_DELAY_TIME: data.delayed_time, - ATTR_PLANNED_TIME: data.planned_time, - ATTR_ESTIMATED_TIME: data.estimated_time, - ATTR_ACTUAL_TIME: data.actual_time, - ATTR_OTHER_INFORMATION: data.other_info, - ATTR_DEVIATIONS: data.deviation, - }, + ), + TrafikverketSensorEntityDescription( + key="departure_state", + translation_key="departure_state", + icon="mdi:clock", + value_fn=lambda data: data.departure_time, + device_class=SensorDeviceClass.ENUM, + options=["on_time", "delayed", "canceled"], + ), + TrafikverketSensorEntityDescription( + key="cancelled", + translation_key="cancelled", + icon="mdi:alert", + value_fn=lambda data: data.cancelled, + ), + TrafikverketSensorEntityDescription( + key="delayed_time", + translation_key="delayed_time", + icon="mdi:clock", + device_class=SensorDeviceClass.DURATION, + native_unit_of_measurement=UnitOfTime.SECONDS, + value_fn=lambda data: data.delayed_time, + ), + TrafikverketSensorEntityDescription( + key="planned_time", + translation_key="planned_time", + icon="mdi:clock", + device_class=SensorDeviceClass.TIMESTAMP, + value_fn=lambda data: data.planned_time, + entity_registry_enabled_default=False, + ), + TrafikverketSensorEntityDescription( + key="estimated_time", + translation_key="estimated_time", + icon="mdi:clock", + device_class=SensorDeviceClass.TIMESTAMP, + value_fn=lambda data: data.estimated_time, + entity_registry_enabled_default=False, + ), + TrafikverketSensorEntityDescription( + key="actual_time", + translation_key="actual_time", + icon="mdi:clock", + device_class=SensorDeviceClass.TIMESTAMP, + value_fn=lambda data: data.actual_time, + entity_registry_enabled_default=False, + ), + TrafikverketSensorEntityDescription( + key="other_info", + translation_key="other_info", + icon="mdi:information-variant", + value_fn=lambda data: data.other_info, + ), + TrafikverketSensorEntityDescription( + key="deviation", + translation_key="deviation", + icon="mdi:alert", + value_fn=lambda data: data.deviation, ), ) @@ -81,71 +113,48 @@ async def async_setup_entry( coordinator: TVDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] - to_station = coordinator.to_station - from_station = coordinator.from_station - get_time: str | None = entry.data.get(CONF_TIME) - train_time = dt_util.parse_time(get_time) if get_time else None - async_add_entities( [ - TrainSensor( - coordinator, - entry.data[CONF_NAME], - from_station, - to_station, - entry.data[CONF_WEEKDAY], - train_time, - entry.entry_id, - description, - ) + TrainSensor(coordinator, entry.data[CONF_NAME], entry.entry_id, description) for description in SENSOR_TYPES - ], - True, + ] ) class TrainSensor(CoordinatorEntity[TVDataUpdateCoordinator], SensorEntity): """Contains data about a train depature.""" - _attr_has_entity_name = True entity_description: TrafikverketSensorEntityDescription + _attr_attribution = ATTRIBUTION + _attr_has_entity_name = True def __init__( self, coordinator: TVDataUpdateCoordinator, name: str, - from_station: StationInfo, - to_station: StationInfo, - weekday: list, - departuretime: time | None, entry_id: str, entity_description: TrafikverketSensorEntityDescription, ) -> None: """Initialize the sensor.""" super().__init__(coordinator) + self._attr_unique_id = f"{entry_id}-{entity_description.key}" self.entity_description = entity_description self._attr_device_info = DeviceInfo( entry_type=DeviceEntryType.SERVICE, identifiers={(DOMAIN, entry_id)}, - manufacturer="Trafikverket", - model="v2.0", name=name, configuration_url="https://api.trafikinfo.trafikverket.se/", ) - self._attr_unique_id = f"{entry_id}-{entity_description.key}" self._update_attr() + @callback + def _update_attr(self) -> None: + """Update _attr.""" + self._attr_native_value = self.entity_description.value_fn( + self.coordinator.data + ) + @callback def _handle_coordinator_update(self) -> None: self._update_attr() return super()._handle_coordinator_update() - - @callback - def _update_attr(self) -> None: - """Retrieve latest states.""" - self._attr_native_value = self.entity_description.value_fn( - self.coordinator.data - ) - self._attr_extra_state_attributes = self.entity_description.extra_fn( - self.coordinator.data - ) diff --git a/homeassistant/components/trafikverket_train/strings.json b/homeassistant/components/trafikverket_train/strings.json index 05032027b97..59431107ae2 100644 --- a/homeassistant/components/trafikverket_train/strings.json +++ b/homeassistant/components/trafikverket_train/strings.json @@ -45,38 +45,36 @@ "entity": { "sensor": { "departure_time": { - "name": "Departure time", - "state_attributes": { - "departure_state": { - "name": "Departure state", - "state": { - "on_time": "On time", - "delayed": "Delayed", - "canceled": "Cancelled" - } - }, - "canceled": { - "name": "Cancelled" - }, - "number_of_minutes_delayed": { - "name": "Minutes delayed" - }, - "planned_time": { - "name": "Planned time" - }, - "estimated_time": { - "name": "Estimated time" - }, - "actual_time": { - "name": "Actual time" - }, - "other_information": { - "name": "Other information" - }, - "deviations": { - "name": "Deviations" - } + "name": "Departure time" + }, + "departure_state": { + "name": "Departure state", + "state": { + "on_time": "On time", + "delayed": "Delayed", + "canceled": "Cancelled" } + }, + "cancelled": { + "name": "Cancelled" + }, + "delayed_time": { + "name": "Delayed time" + }, + "planned_time": { + "name": "Planned time" + }, + "estimated_time": { + "name": "Estimated time" + }, + "actual_time": { + "name": "Actual time" + }, + "other_info": { + "name": "Other information" + }, + "deviation": { + "name": "Deviation" } } }