Add a coordinator to Waze Travel Time (#148585)

Author: etiennec78

homeassistant/components/waze_travel_time/__init__.py

@@ -1,15 +1,13 @@
 """The waze_travel_time component."""
 
 import asyncio
-from collections.abc import Collection
 import logging
-from typing import Literal
 
-from pywaze.route_calculator import CalcRoutesResponse, WazeRouteCalculator, WRCError
+from pywaze.route_calculator import WazeRouteCalculator
 import voluptuous as vol
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_REGION, Platform, UnitOfLength
+from homeassistant.const import CONF_REGION, Platform
 from homeassistant.core import (
     HomeAssistant,
     ServiceCall,
@@ -27,7 +25,6 @@
     TextSelectorConfig,
     TextSelectorType,
 )
-from homeassistant.util.unit_conversion import DistanceConverter
 
 from .const import (
     CONF_AVOID_FERRIES,
@@ -43,13 +40,13 @@
     DEFAULT_FILTER,
     DEFAULT_VEHICLE_TYPE,
     DOMAIN,
-    IMPERIAL_UNITS,
     METRIC_UNITS,
     REGIONS,
     SEMAPHORE,
     UNITS,
     VEHICLE_TYPES,
 )
+from .coordinator import WazeTravelTimeCoordinator, async_get_travel_times
 
 PLATFORMS = [Platform.SENSOR]
 
@@ -109,6 +106,16 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
     if SEMAPHORE not in hass.data.setdefault(DOMAIN, {}):
         hass.data.setdefault(DOMAIN, {})[SEMAPHORE] = asyncio.Semaphore(1)
 
+    httpx_client = get_async_client(hass)
+    client = WazeRouteCalculator(
+        region=config_entry.data[CONF_REGION].upper(), client=httpx_client
+    )
+
+    coordinator = WazeTravelTimeCoordinator(hass, config_entry, client)
+    config_entry.runtime_data = coordinator
+
+    await coordinator.async_config_entry_first_refresh()
+
     await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
 
     async def async_get_travel_times_service(service: ServiceCall) -> ServiceResponse:
@@ -140,7 +147,7 @@ async def async_get_travel_times_service(service: ServiceCall) -> ServiceRespons
             incl_filters=service.data.get(CONF_INCL_FILTER, DEFAULT_FILTER),
             excl_filters=service.data.get(CONF_EXCL_FILTER, DEFAULT_FILTER),
         )
-        return {"routes": [vars(route) for route in response]} if response else None
+        return {"routes": [vars(route) for route in response]}
 
     hass.services.async_register(
         DOMAIN,
@@ -152,106 +159,6 @@ async def async_get_travel_times_service(service: ServiceCall) -> ServiceRespons
     return True
 
 
-async def async_get_travel_times(
-    client: WazeRouteCalculator,
-    origin: str,
-    destination: str,
-    vehicle_type: str,
-    avoid_toll_roads: bool,
-    avoid_subscription_roads: bool,
-    avoid_ferries: bool,
-    realtime: bool,
-    units: Literal["metric", "imperial"] = "metric",
-    incl_filters: Collection[str] | None = None,
-    excl_filters: Collection[str] | None = None,
-) -> list[CalcRoutesResponse] | None:
-    """Get all available routes."""
-
-    incl_filters = incl_filters or ()
-    excl_filters = excl_filters or ()
-
-    _LOGGER.debug(
-        "Getting update for origin: %s destination: %s",
-        origin,
-        destination,
-    )
-    routes = []
-    vehicle_type = "" if vehicle_type.upper() == "CAR" else vehicle_type.upper()
-    try:
-        routes = await client.calc_routes(
-            origin,
-            destination,
-            vehicle_type=vehicle_type,
-            avoid_toll_roads=avoid_toll_roads,
-            avoid_subscription_roads=avoid_subscription_roads,
-            avoid_ferries=avoid_ferries,
-            real_time=realtime,
-            alternatives=3,
-        )
-        _LOGGER.debug("Got routes: %s", routes)
-
-        incl_routes: list[CalcRoutesResponse] = []
-
-        def should_include_route(route: CalcRoutesResponse) -> bool:
-            if len(incl_filters) < 1:
-                return True
-            should_include = any(
-                street_name in incl_filters or "" in incl_filters
-                for street_name in route.street_names
-            )
-            if not should_include:
-                _LOGGER.debug(
-                    "Excluding route [%s], because no inclusive filter matched any streetname",
-                    route.name,
-                )
-                return False
-            return True
-
-        incl_routes = [route for route in routes if should_include_route(route)]
-
-        filtered_routes: list[CalcRoutesResponse] = []
-
-        def should_exclude_route(route: CalcRoutesResponse) -> bool:
-            for street_name in route.street_names:
-                for excl_filter in excl_filters:
-                    if excl_filter == street_name:
-                        _LOGGER.debug(
-                            "Excluding route, because exclusive filter [%s] matched streetname: %s",
-                            excl_filter,
-                            route.name,
-                        )
-                        return True
-            return False
-
-        filtered_routes = [
-            route for route in incl_routes if not should_exclude_route(route)
-        ]
-
-        if units == IMPERIAL_UNITS:
-            filtered_routes = [
-                CalcRoutesResponse(
-                    name=route.name,
-                    distance=DistanceConverter.convert(
-                        route.distance, UnitOfLength.KILOMETERS, UnitOfLength.MILES
-                    ),
-                    duration=route.duration,
-                    street_names=route.street_names,
-                )
-                for route in filtered_routes
-                if route.distance is not None
-            ]
-
-        if len(filtered_routes) < 1:
-            _LOGGER.warning("No routes found")
-            return None
-    except WRCError as exp:
-        _LOGGER.warning("Error on retrieving data: %s", exp)
-        return None
-
-    else:
-        return filtered_routes
-
-
 async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
     """Unload a config entry."""
     return await hass.config_entries.async_unload_platforms(config_entry, PLATFORMS)