Merge pull request #25 from jirutka/typing-and-l10n

Improve typing, localization, flatten attributes etc.

custom_components/pid_departures/__init__.py

@@ -2,20 +2,21 @@
 from __future__ import annotations
 
 from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_API_KEY, CONF_ID
 from homeassistant.core import HomeAssistant
 
-from . import hub
 from .const import DOMAIN, CONF_DEP_NUM
-from homeassistant.const import CONF_API_KEY, CONF_ID
-from .dep_board_api import PIDDepartureBoardAPI
+from .hub import DepartureBoard
 
 PLATFORMS: list[str] = ["sensor", "binary_sensor", "calendar"]
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Departure Board from a config entry flow."""
-    response = await PIDDepartureBoardAPI.async_fetch_data(entry.data[CONF_API_KEY], entry.data[CONF_ID], entry.data[CONF_DEP_NUM])
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = hub.DepartureBoard(hass, entry.data[CONF_API_KEY], entry.data[CONF_ID], entry.data[CONF_DEP_NUM], response)
+    hub = DepartureBoard(hass, entry.data[CONF_API_KEY], entry.data[CONF_ID], entry.data[CONF_DEP_NUM])  # type: ignore[Any]
+    await hub.async_update()
+
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = hub  # type: ignore[Any]
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
     return True
@@ -28,6 +29,6 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     # details
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
     if unload_ok:
-        hass.data[DOMAIN].pop(entry.entry_id)
+        hass.data[DOMAIN].pop(entry.entry_id)  # type: ignore[Any]
 
     return unload_ok

custom_components/pid_departures/binary_sensor.py

@@ -1,45 +1,43 @@
 """Platform for binary sensor."""
 from __future__ import annotations
 
-from homeassistant.components.binary_sensor import BinarySensorEntity, BinarySensorDeviceClass
-from .const import ICON_INFO_ON, DOMAIN, ICON_INFO_OFF, ICON_WHEEL
-from homeassistant.const import EntityCategory, STATE_UNKNOWN, STATE_UNAVAILABLE, STATE_ON, STATE_OFF
+from collections.abc import Mapping
+from typing import Any
 
+from homeassistant.components.binary_sensor import BinarySensorEntity, BinarySensorDeviceClass
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import EntityCategory
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-async def async_setup_entry(hass, config_entry, async_add_entities):
+from .const import ICON_INFO_ON, DOMAIN, ICON_INFO_OFF, ICON_WHEEL
+from .entity import BaseEntity
+from .hub import DepartureBoard
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback
+) -> None:
     """Add sensors for passed config_entry in HA."""
-    departure_board = hass.data[DOMAIN][config_entry.entry_id]
+    departure_board: DepartureBoard = hass.data[DOMAIN][config_entry.entry_id]  # type: ignore[Any]
     async_add_entities([WheelchairSensor(departure_board), InfotextBinarySensor(departure_board)])
 
 
-class InfotextBinarySensor(BinarySensorEntity):
+class InfotextBinarySensor(BaseEntity, BinarySensorEntity):
     """Sensor for info text."""
-    _attr_has_entity_name = True
+
+    _attr_translation_key = "infotext"
     _attr_device_class = BinarySensorDeviceClass.PROBLEM
     _attr_entity_category = EntityCategory.DIAGNOSTIC
     _attr_should_poll = False
 
-    def __init__(self, departure_board):
-
-        self._departure_board = departure_board
-        self._attr_unique_id = f"{self._departure_board.board_id}_{self._departure_board.conn_num+7}"
-
-    @property
-    def device_info(self):
-        """Returns information to link this entity with the correct device."""
-        return self._departure_board.device_info
-
-    @property
-    def name(self) -> str:
-        """Returns entity name"""
-        return "infotext"
-
     @property
     def is_on(self) -> bool | None:
         return self._departure_board.info_text[0]
 
     @property
-    def extra_state_attributes(self):
+    def extra_state_attributes(self) -> Mapping[str, Any]:
         return self._departure_board.info_text[1]
 
     @property
@@ -49,44 +47,30 @@ def icon(self) -> str:
         else:
             return ICON_INFO_OFF
 
-    async def async_added_to_hass(self):
+    async def async_added_to_hass(self) -> None:
         """Run when this Entity has been added to HA."""
         # Sensors should also register callbacks to HA when their state changes
         self._departure_board.register_callback(self.async_write_ha_state)
 
-    async def async_will_remove_from_hass(self):
+    async def async_will_remove_from_hass(self) -> None:
         """Entity being removed from hass."""
         # The opposite of async_added_to_hass. Remove any registered call backs here.
         self._departure_board.remove_callback(self.async_write_ha_state)
 
 
-class WheelchairSensor(BinarySensorEntity):
+class WheelchairSensor(BaseEntity, BinarySensorEntity):
     """Sensor for wheelchair accessibility of the station."""
-    _attr_has_entity_name = True
+
+    _attr_translation_key = "wheelchair_accessible"
     _attr_icon = ICON_WHEEL
     _attr_entity_category = EntityCategory.DIAGNOSTIC
     _attr_should_poll = False
 
-    def __init__(self, departure_board):
-
-        self._departure_board = departure_board
-        self._attr_unique_id = f"{self._departure_board.board_id}_{self._departure_board.conn_num+8}"
-
     @property
-    def device_info(self):
-        """Returns information to link this entity with the correct device."""
-        return self._departure_board.device_info
-
-    @property
-    def name(self):
-        """Returns entity name"""
-        return "wheelchair"
-
-    @property
-    def is_on(self):
+    def is_on(self) -> bool | None:
         if self._departure_board.wheelchair_accessible == 1:
-            return STATE_ON
+            return True
         elif self._departure_board.wheelchair_accessible == 2:
-            return STATE_OFF
+            return False
         else:
             return None

custom_components/pid_departures/calendar.py

@@ -4,19 +4,20 @@
 from collections.abc import Mapping
 from datetime import datetime, timedelta
 import logging
-from typing import Any
+from typing import Any, cast
 from typing_extensions import override
 
 from homeassistant.components.calendar import CalendarEntity, CalendarEvent
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
+from homeassistant.const import CONF_LATITUDE, CONF_LONGITUDE, STATE_ON
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import dt
 
-from custom_components.pid_departures.dep_board_api import PIDDepartureBoardAPI
-
-from .const import CAL_EVENT_MIN_DURATION_SEC, CONF_CAL_EVENTS_NUM, DOMAIN
-from .hub import DepartureBoard
+from .const import CAL_EVENT_MIN_DURATION_SEC, CONF_CAL_EVENTS_NUM, DOMAIN, ICON_STOP, ROUTE_TYPE_ICON, RouteType
+from .dep_board_api import PIDDepartureBoardAPI
+from .entity import BaseEntity
+from .hub import DepartureBoard, DepartureData
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -33,18 +34,13 @@ async def async_setup_entry(
     ])
 
 
-class DeparturesCalendarEntity(CalendarEntity):
-    _attr_has_entity_name = True
+class DeparturesCalendarEntity(BaseEntity, CalendarEntity):
+
     _attr_should_poll = False
     _attr_translation_key = "departures"
 
     def __init__(self, departure_board: DepartureBoard, events_count: int) -> None:
-        super().__init__()
-        self._attr_unique_id = f"{departure_board.board_id}_{departure_board.conn_num}"
-        self._attr_translation_placeholders = {
-            "stop_name": departure_board.name,
-        }
-        self._departure_board = departure_board
+        super().__init__(departure_board)
         self._events_count = events_count
         self._event: CalendarEvent | None = None
 
@@ -64,12 +60,28 @@ async def async_will_remove_from_hass(self):
     @override
     def event(self) -> CalendarEvent | None:
         """Return the current or next upcoming event."""
-        return self._create_event(self._departure_board.extra_attr[0])
+        return self._create_event(self._departure_board.departures[0])
+
+    @property
+    @override
+    def icon(self) -> str:
+        """Return entity icon based on the type of route."""
+        if self.state == STATE_ON:
+            route_type = self._departure_board.departures[0].route_type
+            return ROUTE_TYPE_ICON.get(route_type, ROUTE_TYPE_ICON[RouteType.BUS])
+        else:
+            return ICON_STOP
 
     @property
     @override
     def extra_state_attributes(self) -> Mapping[str, Any]:
-        return self._departure_board.extra_attr[0]
+        # NOTE: When CONF_LATITUDE and CONF_LONGITUDE is included, HASS shows
+        #  the entity on the map.
+        return {
+            **self._departure_board.departures[0].as_dict(),
+            CONF_LATITUDE: self._departure_board.latitude,
+            CONF_LONGITUDE: self._departure_board.longitude,
+        }
 
     @override
     async def async_get_events(
@@ -92,12 +104,15 @@ async def async_get_events(
             time_before=timedelta_clamp(time_before, *PIDDepartureBoardAPI.TIME_BEFORE_RANGE),
             time_after=timedelta_clamp(time_after, *PIDDepartureBoardAPI.TIME_AFTER_RANGE))
 
-        events = (self._create_event(dep) for dep in data["departures"])
+        events = (
+            self._create_event(DepartureData.from_api(dep))
+            for dep in cast(list[dict[str, Any]], data["departures"])
+        )
         return [event for event in events if event]
 
-    def _create_event(self, departure: dict[str, Any]) -> CalendarEvent | None:
-        start = try_parse_timestamp(departure["arrival_timestamp"]["predicted"])  # type: ignore[Any]
-        end = try_parse_timestamp(departure["departure_timestamp"]["predicted"])  # type: ignore[Any]
+    def _create_event(self, departure: DepartureData) -> CalendarEvent | None:
+        start = departure.arrival_time_est
+        end = departure.departure_time_est
 
         if not start and not end:
             _LOGGER.error('Invalid data, both "arrival_timestamp" and "departure_timestamp" is null')
@@ -110,29 +125,22 @@ def _create_event(self, departure: dict[str, Any]) -> CalendarEvent | None:
             # arrival_timestamp is null on first stops.
             start = end - timedelta(seconds=CAL_EVENT_MIN_DURATION_SEC)
 
-        route_type_code: int = departure["route"]["type"]
-        route_type_name = self._translate(f"state_attributes.route_type.state.{route_type_code}")
-        short_name: str = departure["route"]["short_name"]
+        route_type = self._translate(f"state_attributes.route_type.state.{departure.route_type}")
+        short_name = departure.route_name or "?"
 
         return CalendarEvent(
             start=start,
             end=end,
-            summary=f"{route_type_name} {short_name}",
+            summary=f"{route_type} {short_name}",
             location=self._departure_board.name,
+            description=f"Trip to {departure.trip_headsign}",
         )
 
-    def _translate(self, key_path: str) -> str | None:
+    def _translate(self, key_path: str) -> str:
         # XXX: This is hack-ish, I haven't found the right approach for this.
-        return self.platform.platform_translations.get(
+        return self.platform.platform_translations[
             f"component.{self.platform.platform_name}.entity.{self.platform.domain}" +
-            f".{self.translation_key}.{key_path}")
-
-
-def try_parse_timestamp(input: str | None) -> datetime | None:
-    try:
-        return datetime.fromisoformat(input or "")
-    except ValueError:
-        return None
+            f".{self.translation_key}.{key_path}"]
 
 
 def timedelta_clamp(delta: timedelta, min: timedelta, max: timedelta) -> timedelta:

custom_components/pid_departures/config_flow.py

@@ -1,30 +1,31 @@
-import voluptuous as vol
-
 import logging
-from typing import Any, Tuple, Dict
-from .dep_board_api import PIDDepartureBoardAPI
+from typing import Any, cast
 
-from .const import CONF_CAL_EVENTS_NUM, CONF_DEP_NUM, CONF_STOP_SEL, DOMAIN
 from homeassistant.const import CONF_API_KEY, CONF_ID
 from homeassistant import config_entries
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.selector import selector
-from .stop_list import STOP_LIST, ASW_IDS
+import voluptuous as vol
+
+from .const import CONF_CAL_EVENTS_NUM, CONF_DEP_NUM, CONF_STOP_SEL, DOMAIN
+from .dep_board_api import PIDDepartureBoardAPI
 from .errors import CannotConnect, NoDeparturesSelected, StopNotFound, StopNotInList, WrongApiKey
+from .hub import DepartureBoard
+from .stop_list import STOP_LIST, ASW_IDS
 
 _LOGGER = logging.getLogger(__name__)
 
 
-async def validate_input(hass: HomeAssistant, data: dict) -> tuple[dict[str, str], dict]:
+async def validate_input(hass: HomeAssistant, data: dict[str, Any]) -> tuple[dict[str, str], dict[str, Any]]:
     """Validate the user input allows us to connect.
     Data has the keys from DATA_SCHEMA with values provided by the user.
     """
     try:
-        data[CONF_ID] = ASW_IDS[STOP_LIST.index(data[CONF_STOP_SEL])]
+        data[CONF_ID] = ASW_IDS[STOP_LIST.index(data[CONF_STOP_SEL])]  # type: ignore[Any]
     except Exception:
         raise StopNotInList
 
-    reply = await PIDDepartureBoardAPI.async_fetch_data(data[CONF_API_KEY], data[CONF_ID], data[CONF_DEP_NUM])
+    reply = await PIDDepartureBoardAPI.async_fetch_data(data[CONF_API_KEY], data[CONF_ID], data[CONF_DEP_NUM])  # type: ignore[Any]
 
     title: str = reply["stops"][0]["stop_name"] + " " + (reply["stops"][0]["platform_code"] or "")
     if data[CONF_DEP_NUM] == 0:
@@ -38,17 +39,16 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     CONNECTION_CLASS = config_entries.CONN_CLASS_LOCAL_POLL
     VERSION = 0.1
 
-    async def async_step_user(self, user_input=None):
-
+    async def async_step_user(self, user_input: dict[str, Any] | None = None) -> config_entries.FlowResult:
         # Check for any previous instance of the integration
-        if DOMAIN in list(self.hass.data.keys()):
-            # If previous instance exists, set the API key as suggestion to new config
-            data_schema = {vol.Required(CONF_API_KEY, default=self.hass.data[DOMAIN][list(self.hass.data[DOMAIN].keys())[0]].api_key): str}
-        else:
-            # if no previous instance, show blank form
-            data_schema = {vol.Required(CONF_API_KEY): str}
-
-        data_schema.update({
+        api_key: str | None = None
+        if (boards := self.hass.data.get(DOMAIN, {})):  # type: ignore[Any]
+            board: DepartureBoard = next(iter(boards.values()))  # type: ignore[Any]
+            # If previous instance exists, use the API key as suggestion to new config
+            api_key = board.api_key
+
+        data_schema: dict[Any, Any] = {
+            vol.Required(CONF_API_KEY, default=api_key): str,
             vol.Required(CONF_DEP_NUM, default=1): int,
             CONF_STOP_SEL: selector({
                 "select": {
@@ -62,10 +62,10 @@ async def async_step_user(self, user_input=None):
                 vol.Coerce(int),
                 vol.Range(0, 1000),
             ),
-        })
+        }
 
         # Set dict for errors
-        errors: dict = {}
+        errors: dict[str, str] = {}
 
         # Steps to take if user input is received
         if user_input is not None: