coordinator.py

Go to the documentation of this file.

"""NextBus data update coordinator."""
 
from datetime import timedelta
import logging
from typing import Any
 
from py_nextbus import NextBusClient
from py_nextbus.client import NextBusFormatError, NextBusHTTPError
 
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 NextBusDataUpdateCoordinator(DataUpdateCoordinator):
    """Class to manage fetching NextBus data."""
 
    def __init__(self, hass: HomeAssistant, agency: str) -> None:
        """Initialize a global coordinator for fetching data for a given agency."""
        super().__init__(
            hass,
            _LOGGER,
            config_entry=None,  # It is shared between multiple entries
            name=DOMAIN,
            update_interval=timedelta(seconds=30),
        )
        self.clientclient = NextBusClient(agency_id=agency)
        self._agency_agency = agency
        self._route_stops: set[RouteStop] = set()
        self._predictions_predictions: dict[RouteStop, dict[str, Any]] = {}
 
    def add_stop_route(self, stop_id: str, route_id: str) -> None:
        """Tell coordinator to start tracking a given stop and route."""
        self._route_stops.add(RouteStop(route_id, stop_id))
 
    def remove_stop_route(self, stop_id: str, route_id: str) -> None:
        """Tell coordinator to stop tracking a given stop and route."""
        self._route_stops.remove(RouteStop(route_id, stop_id))
 
    def get_prediction_data(self, stop_id: str, route_id: str) -> dict[str, Any] | None:
        """Get prediction result for a given stop and route."""
        return self._predictions_predictions.get(RouteStop(route_id, stop_id))
 
    def has_routes(self) -> bool:
        """Check if this coordinator is tracking any routes."""
        return len(self._route_stops) > 0
 
    async def _async_update_data(self) -> dict[str, Any]:
        """Fetch data from NextBus."""
 
        _stops_to_route_stops: dict[str, set[RouteStop]] = {}
        for route_stop in self._route_stops:
            _stops_to_route_stops.setdefault(route_stop.stop_id, set()).add(route_stop)
 
        self.loggerlogger.debug(
            "Updating data from API. Routes: %s", str(_stops_to_route_stops)
        )
 
        def _update_data() -> dict:
            """Fetch data from NextBus."""
            self.loggerlogger.debug("Updating data from API (executor)")
            predictions: dict[RouteStop, dict[str, Any]] = {}
 
            for stop_id, route_stops in _stops_to_route_stops.items():
                self.loggerlogger.debug("Updating data from API (executor) %s", stop_id)
                try:
                    prediction_results = self.clientclient.predictions_for_stop(stop_id)
                except NextBusHTTPError as ex:
                    self.loggerlogger.error(
                        "Error updating %s (executor): %s %s",
                        str(stop_id),
                        ex,
                        getattr(ex, "response", None),
                    )
                    raise UpdateFailed("Failed updating nextbus data", ex) from ex
                except NextBusFormatError as ex:
                    raise UpdateFailed("Failed updating nextbus data", ex) from ex
 
                self.loggerlogger.debug(
                    "Prediction results for %s (executor): %s",
                    str(stop_id),
                    str(prediction_results),
                )
 
                for route_stop in route_stops:
                    for prediction_result in prediction_results:
                        if (
                            prediction_result["stop"]["id"] == route_stop.stop_id
                            and prediction_result["route"]["id"] == route_stop.route_id
                        ):
                            predictions[route_stop] = prediction_result
                            break
                    else:
                        self.loggerlogger.warning(
                            "Prediction not found for %s (executor)", str(route_stop)
                        )
 
            self._predictions_predictions = predictions
 
            return predictions
 
        return await self.hasshass.async_add_executor_job(_update_data)