Add entity support to Waze Travel Time (#14934)

Current version only supports latitude and longitude or an address for the origin and destination fields. This update allows those fields to use entity IDs of device_tracker, zone, and sensor.

homeassistant/components/sensor/waze_travel_time.py

@@ -7,12 +7,14 @@ https://home-assistant.io/components/sensor.waze_travel_time/
 from datetime import timedelta
 import logging
 
-import requests
 import voluptuous as vol
 
 from homeassistant.components.sensor import PLATFORM_SCHEMA
-from homeassistant.const import ATTR_ATTRIBUTION, CONF_NAME, CONF_REGION
+from homeassistant.const import (
+    ATTR_ATTRIBUTION, CONF_NAME, CONF_REGION, EVENT_HOMEASSISTANT_START,
+    ATTR_LATITUDE, ATTR_LONGITUDE)
 import homeassistant.helpers.config_validation as cv
+import homeassistant.helpers.location as location
 from homeassistant.helpers.entity import Entity
 from homeassistant.util import Throttle
 
@@ -20,6 +22,7 @@ REQUIREMENTS = ['WazeRouteCalculator==0.5']
 
 _LOGGER = logging.getLogger(__name__)
 
+ATTR_DURATION = 'duration'
 ATTR_DISTANCE = 'distance'
 ATTR_ROUTE = 'route'
 
@@ -46,6 +49,8 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
     vol.Optional(CONF_EXCL_FILTER): cv.string,
 })
 
+TRACKABLE_DOMAINS = ['device_tracker', 'sensor', 'zone']
+
 
 def setup_platform(hass, config, add_devices, discovery_info=None):
     """Set up the Waze travel time sensor platform."""
@@ -56,24 +61,46 @@ def setup_platform(hass, config, add_devices, discovery_info=None):
     incl_filter = config.get(CONF_INCL_FILTER)
     excl_filter = config.get(CONF_EXCL_FILTER)
 
-    try:
+    sensor = WazeTravelTime(name, origin, destination, region,
-        waze_data = WazeRouteData(
+                            incl_filter, excl_filter)
-            origin, destination, region, incl_filter, excl_filter)
-    except requests.exceptions.HTTPError as error:
-        _LOGGER.error("%s", error)
-        return
 
-    add_devices([WazeTravelTime(waze_data, name)], True)
+    add_devices([sensor], True)
+
+    hass.bus.listen_once(EVENT_HOMEASSISTANT_START, sensor.update)
+
+
+def _get_location_from_attributes(state):
+    """Get the lat/long string from an states attributes."""
+    attr = state.attributes
+    return '{},{}'.format(
+        attr.get(ATTR_LATITUDE),
+        attr.get(ATTR_LONGITUDE)
+    )
 
 
 class WazeTravelTime(Entity):
     """Representation of a Waze travel time sensor."""
 
-    def __init__(self, waze_data, name):
+    def __init__(self, name, origin, destination, region,
+                 incl_filter, excl_filter):
         """Initialize the Waze travel time sensor."""
         self._name = name
+        self._region = region
+        self._incl_filter = incl_filter
+        self._excl_filter = excl_filter
         self._state = None
-        self.waze_data = waze_data
+        self._origin_entity_id = None
+        self._destination_entity_id = None
+
+        if origin.split('.', 1)[0] in TRACKABLE_DOMAINS:
+            self._origin_entity_id = origin
+        else:
+            self._origin = origin
+
+        if destination.split('.', 1)[0] in TRACKABLE_DOMAINS:
+            self._destination_entity_id = destination
+        else:
+            self._destination = destination
 
     @property
     def name(self):
@@ -83,7 +110,12 @@ class WazeTravelTime(Entity):
     @property
     def state(self):
         """Return the state of the sensor."""
-        return round(self._state['duration'])
+        if self._state is None:
+            return None
+
+        if 'duration' in self._state:
+            return round(self._state['duration'])
+        return None
 
     @property
     def unit_of_measurement(self):
@@ -98,54 +130,97 @@ class WazeTravelTime(Entity):
     @property
     def device_state_attributes(self):
         """Return the state attributes of the last update."""
-        return {
+        if self._state is None:
-            ATTR_ATTRIBUTION: CONF_ATTRIBUTION,
+            return None
-            ATTR_DISTANCE: round(self._state['distance']),
-            ATTR_ROUTE: self._state['route'],
-        }
 
-    def update(self):
+        res = {ATTR_ATTRIBUTION: CONF_ATTRIBUTION}
-        """Fetch new state data for the sensor."""
+        if 'duration' in self._state:
-        try:
+            res[ATTR_DURATION] = self._state['duration']
-            self.waze_data.update()
+        if 'distance' in self._state:
-            self._state = self.waze_data.data
+            res[ATTR_DISTANCE] = self._state['distance']
-        except KeyError:
+        if 'route' in self._state:
-            _LOGGER.error("Error retrieving data from server")
+            res[ATTR_ROUTE] = self._state['route']
+        return res
 
+    def _get_location_from_entity(self, entity_id):
+        """Get the location from the entity_id."""
+        state = self.hass.states.get(entity_id)
 
-class WazeRouteData(object):
+        if state is None:
-    """Get data from Waze."""
+            _LOGGER.error("Unable to find entity %s", entity_id)
+            return None
 
-    def __init__(self, origin, destination, region, incl_filter, excl_filter):
+        # Check if the entity has location attributes (zone)
-        """Initialize the data object."""
+        if location.has_location(state):
-        self._destination = destination
+            return _get_location_from_attributes(state)
-        self._origin = origin
+
-        self._region = region
+        # Check if device is in a zone (device_tracker)
-        self._incl_filter = incl_filter
+        zone_state = self.hass.states.get('zone.{}'.format(state.state))
-        self._excl_filter = excl_filter
+        if location.has_location(zone_state):
-        self.data = {}
+            _LOGGER.debug(
+                "%s is in %s, getting zone location",
+                entity_id, zone_state.entity_id
+            )
+            return _get_location_from_attributes(zone_state)
+
+        # If zone was not found in state then use the state as the location
+        if entity_id.startswith('sensor.'):
+            return state.state
+
+        # When everything fails just return nothing
+        return None
+
+    def _resolve_zone(self, friendly_name):
+        """Get a lat/long from a zones friendly_name."""
+        states = self.hass.states.all()
+        for state in states:
+            if state.domain == 'zone' and state.name == friendly_name:
+                return _get_location_from_attributes(state)
+
+        return friendly_name
 
     @Throttle(SCAN_INTERVAL)
     def update(self):
-        """Fetch latest data from Waze."""
+        """Fetch new state data for the sensor."""
         import WazeRouteCalculator
-        _LOGGER.debug("Update in progress...")
+
-        try:
+        if self._origin_entity_id is not None:
-            params = WazeRouteCalculator.WazeRouteCalculator(
+            self._origin = self._get_location_from_entity(
-                self._origin, self._destination, self._region, None)
+                self._origin_entity_id
-            results = params.calc_all_routes_info()
+            )
-            if self._incl_filter is not None:
+
-                results = {k: v for k, v in results.items() if
+        if self._destination_entity_id is not None:
-                           self._incl_filter.lower() in k.lower()}
+            self._destination = self._get_location_from_entity(
-            if self._excl_filter is not None:
+                self._destination_entity_id
-                results = {k: v for k, v in results.items() if
+            )
-                           self._excl_filter.lower() not in k.lower()}
+
-            best_route = next(iter(results))
+        self._destination = self._resolve_zone(self._destination)
-            (duration, distance) = results[best_route]
+        self._origin = self._resolve_zone(self._origin)
-            best_route_str = bytes(best_route, 'ISO-8859-1').decode('UTF-8')
+
-            self.data['duration'] = duration
+        if self._destination is not None and self._origin is not None:
-            self.data['distance'] = distance
+            try:
-            self.data['route'] = best_route_str
+                params = WazeRouteCalculator.WazeRouteCalculator(
-        except WazeRouteCalculator.WRCError as exp:
+                    self._origin, self._destination, self._region)
-            _LOGGER.error("Error on retrieving data: %s", exp)
+                routes = params.calc_all_routes_info()
-            return
+
+                if self._incl_filter is not None:
+                    routes = {k: v for k, v in routes.items() if
+                              self._incl_filter.lower() in k.lower()}
+
+                if self._excl_filter is not None:
+                    routes = {k: v for k, v in routes.items() if
+                              self._excl_filter.lower() not in k.lower()}
+
+                route = sorted(routes, key=(lambda key: routes[key][0]))[0]
+                duration, distance = routes[route]
+                route = bytes(route, 'ISO-8859-1').decode('UTF-8')
+                self._state = {
+                    'duration': duration,
+                    'distance': distance,
+                    'route': route}
+            except WazeRouteCalculator.WRCError as exp:
+                _LOGGER.error("Error on retrieving data: %s", exp)
+                return
+            except KeyError:
+                _LOGGER.error("Error retrieving data from server")
+                return