Add CityBus integration

homeassistant/components/citybus/__init__.py

@@ -0,0 +1,58 @@
+"""CityBus platform."""
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import Platform
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryNotReady
+
+from .const import CONF_ROUTE, CONF_DIRECTION, CONF_STOP, DOMAIN
+from .coordinator import CityBusDataUpdateCoordinator
+
+PLATFORMS = [Platform.SENSOR]
+
+
+async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+    """Set up platforms for CityBus."""
+    entry_route = entry.data[CONF_ROUTE]
+    entry_direction = entry.data[CONF_DIRECTION]
+    entry_stop = entry.data[CONF_STOP]
+    coordinator_key = f"{entry_route}-{entry_direction}-{entry_stop}"
+
+    coordinator: CityBusDataUpdateCoordinator | None = hass.data.setdefault(
+        DOMAIN, {}
+    ).get(
+        coordinator_key,
+    )
+    if coordinator is None:
+        coordinator = CityBusDataUpdateCoordinator(hass)
+        hass.data[DOMAIN][coordinator_key] = coordinator
+
+    coordinator.add_route_stop(entry_route, entry_direction, entry_stop)
+
+    await coordinator.async_refresh()
+    if not coordinator.last_update_success:
+        raise ConfigEntryNotReady from coordinator.last_exception
+
+    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
+
+    return True
+
+
+async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+    """Unload a config entry."""
+    if await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
+        entry_route = entry.data[CONF_ROUTE]
+        entry_direction = entry.data[CONF_DIRECTION]
+        entry_stop = entry.data[CONF_STOP]
+        coordinator_key = f"{entry_route}-{entry_direction}-{entry_stop}"
+
+        coordinator: CityBusDataUpdateCoordinator = hass.data[DOMAIN][coordinator_key]
+        coordinator.remove_route_stop(entry_route, entry_direction, entry_stop)
+
+        if not coordinator.has_route_stops():
+            await coordinator.async_shutdown()
+            hass.data[DOMAIN].pop(coordinator_key)
+
+        return True
+
+    return False

homeassistant/components/citybus/config_flow.py

@@ -0,0 +1,184 @@
+"""Config flow to configure the CityBus integration."""
+
+import logging
+
+from citybussin import Citybussin
+
+import voluptuous as vol
+
+from homeassistant.config_entries import ConfigFlow, ConfigFlowResult
+from homeassistant.helpers.selector import (
+    SelectOptionDict,
+    SelectSelector,
+    SelectSelectorConfig,
+    SelectSelectorMode,
+)
+
+from .const import CONF_DIRECTION, CONF_ROUTE, CONF_STOP, DOMAIN
+
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def _dict_to_select_name_selector(options: dict[str, str]) -> SelectSelector:
+    return SelectSelector(
+        SelectSelectorConfig(
+            options=sorted(
+                (
+                    SelectOptionDict(value=value, label=value)
+                    for key, value in options.items()
+                ),
+                key=lambda o: o["label"],
+            ),
+            mode=SelectSelectorMode.DROPDOWN,
+        )
+    )
+
+
+def _dict_to_select_selector(options: dict[str, str]) -> SelectSelector:
+    return SelectSelector(
+        SelectSelectorConfig(
+            options=sorted(
+                (
+                    SelectOptionDict(value=key, label=value)
+                    for key, value in options.items()
+                ),
+                key=lambda o: o["label"],
+            ),
+            mode=SelectSelectorMode.DROPDOWN,
+        )
+    )
+
+
+def _get_routes(citybussin: Citybussin) -> dict[str, str]:
+    return {a["key"]: a["shortName"] for a in citybussin.get_bus_routes()}
+
+
+def _get_route_key_from_route_name(citybussin: Citybussin, route_name: str) -> str:
+    return [
+        key for key, value in _get_routes(citybussin).items() if value == route_name
+    ][0]
+
+
+def _get_directions(citybussin: Citybussin, route_key: str) -> dict[str, str]:
+    return {
+        a["direction"]["key"]: a["destination"]
+        for a in citybussin.get_route_directions(route_key)
+    }
+
+
+def _get_stops(citybussin: Citybussin, route_key: str) -> dict[str, str]:
+    return {a["stopCode"]: a["name"] for a in citybussin.get_route_stops(route_key)}
+
+
+def _unique_id_from_data(data: dict[str, str]) -> str:
+    return f"{data[CONF_ROUTE]}_{data[CONF_DIRECTION]}_{data[CONF_STOP]}"
+
+
+class CityBusFlowHandler(ConfigFlow, domain=DOMAIN):
+    """Handle CityBus configuration."""
+
+    VERSION = 1
+
+    _route_key: dict[str, str]
+    _direction_key: dict[str, str]
+    _stop_code: dict[str, str]
+
+    def __init__(self) -> None:
+        """Initialize CityBus config flow."""
+        self.data: dict[str, str] = {}
+        self._citybussin = Citybussin()
+
+    async def async_step_user(
+        self,
+        user_input: dict[str, str] | None = None,
+    ) -> ConfigFlowResult:
+        """Handle a flow initiated by the user."""
+        return await self.async_step_route(user_input)
+
+    async def async_step_route(
+        self,
+        user_input: dict[str, str] | None = None,
+    ) -> ConfigFlowResult:
+        """Select route."""
+        if user_input is not None:
+            self.data[CONF_ROUTE] = user_input[CONF_ROUTE]
+
+            return await self.async_step_direction()
+
+        self._routes = await self.hass.async_add_executor_job(
+            _get_routes, self._citybussin
+        )
+
+        return self.async_show_form(
+            step_id="route",
+            data_schema=vol.Schema(
+                {vol.Required(CONF_ROUTE): _dict_to_select_name_selector(self._routes)}
+            ),
+        )
+
+    async def async_step_direction(
+        self,
+        user_input: dict[str, str] | None = None,
+    ) -> ConfigFlowResult:
+        """Select direction."""
+        if user_input is not None:
+            self.data[CONF_DIRECTION] = user_input[CONF_DIRECTION]
+
+            return await self.async_step_stop()
+
+        self._route_key = await self.hass.async_add_executor_job(
+            _get_route_key_from_route_name, self._citybussin, self.data[CONF_ROUTE]
+        )
+
+        self._directions = await self.hass.async_add_executor_job(
+            _get_directions, self._citybussin, self._route_key
+        )
+
+        return self.async_show_form(
+            step_id="direction",
+            data_schema=vol.Schema(
+                {
+                    vol.Required(CONF_DIRECTION): _dict_to_select_name_selector(
+                        self._directions
+                    )
+                }
+            ),
+        )
+
+    async def async_step_stop(
+        self,
+        user_input: dict[str, str] | None = None,
+    ) -> ConfigFlowResult:
+        """Select stop."""
+        if user_input is not None:
+            self.data[CONF_STOP] = user_input[CONF_STOP]
+
+            await self.async_set_unique_id(_unique_id_from_data(self.data))
+            self._abort_if_unique_id_configured()
+
+            route_name = self.data[CONF_ROUTE]
+            direction_destination = self.data[CONF_DIRECTION]
+
+            self._route_key = await self.hass.async_add_executor_job(
+                _get_route_key_from_route_name, self._citybussin, route_name
+            )
+
+            stop_code = self.data[CONF_STOP]
+            stop_name = self._stops[stop_code]
+
+            return self.async_create_entry(
+                title=f"{route_name} - {direction_destination} - {stop_name}",
+                data=self.data,
+            )
+
+        self._stops = await self.hass.async_add_executor_job(
+            _get_stops, self._citybussin, self._route_key
+        )
+
+        return self.async_show_form(
+            step_id="stop",
+            data_schema=vol.Schema(
+                {vol.Required(CONF_STOP): _dict_to_select_selector(self._stops)}
+            ),
+        )

homeassistant/components/citybus/const.py

@@ -0,0 +1,7 @@
+"""CityBus constants"""
+
+DOMAIN = "citybus"
+
+CONF_ROUTE = "route"
+CONF_DIRECTION = "direction"
+CONF_STOP = "stop"

homeassistant/components/citybus/coordinator.py

@@ -0,0 +1,86 @@
+"""CityBus data update coordinator."""
+
+from datetime import timedelta
+import logging
+from typing import Any
+
+from citybussin import Citybussin
+
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import DOMAIN
+from .util import RouteStop
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class CityBusDataUpdateCoordinator(DataUpdateCoordinator):
+    """Class to manage fetching CityBus data."""
+
+    def __init__(self, hass: HomeAssistant) -> None:
+        """Initialize a global coordinator for fetching data."""
+        super().__init__(
+            hass,
+            _LOGGER,
+            name=DOMAIN,
+            update_interval=timedelta(seconds=30),
+        )
+        self.citybussin = Citybussin()
+        self._route_stops: set[RouteStop] = set()
+        self._estimates: dict[RouteStop, dict[str, Any]] = {}
+
+    def add_route_stop(
+        self, route_name: str, direction_destination: str, stop_code: str
+    ) -> None:
+        """Tell coordinator to start tracking a given stop for a route and direction."""
+        self._route_stops.add(RouteStop(route_name, direction_destination, stop_code))
+
+    def remove_route_stop(
+        self, route_name: str, direction_destination: str, stop_code: str
+    ) -> None:
+        """Tell coordinator to stop tracking a given stop for a route and direction."""
+        self._route_stops.remove(
+            RouteStop(route_name, direction_destination, stop_code)
+        )
+
+    def get_estimate_data(
+        self, route_name: str, direction_destination: str, stop_code: str
+    ) -> dict[str, Any] | None:
+        """Get the estimate data for a given stop for a route and direction."""
+        return self._estimates.get(
+            RouteStop(route_name, direction_destination, stop_code)
+        )
+
+    def has_route_stops(self) -> bool:
+        """Check if this coordinator is tracking any route stops."""
+        return len(self._route_stops) > 0
+
+    async def _async_update_data(self) -> dict[str, Any]:
+        """Fetch data from CityBus."""
+
+        def _update_data() -> dict:
+            """Fetch data from CityBus."""
+            self.logger.debug("Updating data from API (executor)")
+            estimates: dict[RouteStop, dict[str, Any]] = {}
+
+            for route_stop in self._route_stops:
+                try:
+                    route_key = self.citybussin.get_route_by_short_name(
+                        route_stop.route_name
+                    )["key"]
+                    direction_key = self.citybussin.get_direction_by_destination(
+                        route_key, route_stop.direction_destination
+                    )["direction"]["key"]
+                    estimates[route_stop] = self.citybussin.get_next_depart_times(
+                        route_key, direction_key, route_stop.stop_code
+                    )
+                except Exception as err:
+                    raise UpdateFailed(
+                        f"Error fetching data for CityBus stop {route_stop}: {err}"
+                    ) from err
+
+            self._estimates = estimates
+            return estimates
+
+        return await self.hass.async_add_executor_job(_update_data)

homeassistant/components/citybus/icons.json

@@ -0,0 +1,9 @@
+{
+    "entity": {
+      "sensor": {
+        "citybus": {
+          "default": "mdi:bus"
+        }
+      }
+    }
+}

homeassistant/components/citybus/manifest.json

@@ -0,0 +1,10 @@
+{
+  "domain": "citybus",
+  "name": "CityBus",
+  "codeowners": ["@ericswpark"],
+  "config_flow": true,
+  "documentation": "https://www.home-assistant.io/integrations/citybus/",
+  "iot_class": "cloud_polling",
+  "requirements": ["citybussin==0.0.1a3"],
+  "version": "0.0.1"
+}