Add origin/destination attributes

custom_components/here_travel_time/sensor.py

@@ -1,7 +1,8 @@
 """Support for HERE travel time sensors."""
-from datetime import datetime, timedelta
+from datetime import timedelta
 import logging
 import re
+from typing import Callable, Dict, Optional, Union
 
 import herepy
 import voluptuous as vol
@@ -10,10 +11,10 @@
 from homeassistant.const import (
     ATTR_ATTRIBUTION, ATTR_LATITUDE, ATTR_LONGITUDE, CONF_MODE, CONF_NAME,
     CONF_UNIT_SYSTEM, CONF_UNIT_SYSTEM_IMPERIAL, CONF_UNIT_SYSTEM_METRIC)
+from homeassistant.core import HomeAssistant, State
 from homeassistant.helpers import location
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity import Entity
-import homeassistant.util.dt as dt_util
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -53,6 +54,8 @@
 ATTR_DURATION = 'duration'
 ATTR_DISTANCE = 'distance'
 ATTR_ROUTE = 'route'
+ATTR_ORIGIN = 'origin'
+ATTR_DESTINATION = 'destination'
 
 ATTR_DURATION_WITHOUT_TRAFFIC = 'duration_without_traffic'
 ATTR_ORIGIN_NAME = 'origin_name'
@@ -65,6 +68,10 @@
 TRACKABLE_DOMAINS = ['device_tracker', 'sensor', 'zone', 'person']
 DATA_KEY = 'here_travel_time'
 
+NO_ROUTE_ERRORS = [
+    'NGEO_ERROR_GRAPH_DISCONNECTED',
+    'NGEO_ERROR_ROUTE_NO_END_POINT'
+]
 NO_ROUTE_ERROR_MESSAGE = "HERE could not find a route based on the input"
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
@@ -84,18 +91,11 @@
 )
 
 
-def convert_time_to_utc(timestr):
-    """Take a string like 08:00:00 and convert it to a unix timestamp."""
-    combined = datetime.combine(
-        dt_util.start_of_local_day(), dt_util.parse_time(timestr)
-    )
-    if combined < datetime.now():
-        combined = combined + timedelta(days=1)
-    return dt_util.as_timestamp(combined)
-
-
-async def async_setup_platform(hass, config, async_add_entities,
-                               discovery_info=None):
+async def async_setup_platform(
+        hass: HomeAssistant,
+        config: Dict[str, Union[str, bool]],
+        async_add_entities: Callable,
+        discovery_info: None = None) -> None:
     """Set up the HERE travel time platform."""
     hass.data.setdefault(DATA_KEY, [])
 
@@ -133,7 +133,13 @@ async def async_setup_platform(hass, config, async_add_entities,
 class HERETravelTimeSensor(Entity):
     """Representation of a HERE travel time sensor."""
 
-    def __init__(self, hass, name, origin, destination, here_data):
+    def __init__(
+            self,
+            hass: HomeAssistant,
+            name: str,
+            origin: str,
+            destination: str,
+            here_data: 'HERETravelTimeData') -> None:
         """Initialize the sensor."""
         self._hass = hass
         self._name = name
@@ -154,20 +160,21 @@ def __init__(self, hass, name, origin, destination, here_data):
             self._here_data.destination = destination
 
     @property
-    def state(self):
+    def state(self) -> Optional[int]:
         """Return the state of the sensor."""
         if self._here_data.duration is not None:
             return round(self._here_data.duration / 60)
 
         return None
 
     @property
-    def name(self):
+    def name(self) -> str:
         """Get the name of the sensor."""
         return self._name
 
     @property
-    def device_state_attributes(self):
+    def device_state_attributes(
+            self) -> Optional[Dict[str, Union[None, float, str, bool]]]:
         """Return the state attributes."""
         if self._here_data.duration is None:
             return None
@@ -179,19 +186,21 @@ def device_state_attributes(self):
         res[ATTR_ROUTE] = self._here_data.route
         res[CONF_UNIT_SYSTEM] = self._here_data.units
         res[ATTR_DURATION_WITHOUT_TRAFFIC] = self._here_data.base_time / 60
+        res[ATTR_ORIGIN] = self._here_data.origin
+        res[ATTR_DESTINATION] = self._here_data.destination
         res[ATTR_ORIGIN_NAME] = self._here_data.origin_name
         res[ATTR_DESTINATION_NAME] = self._here_data.destination_name
         res[CONF_MODE] = self._here_data.travel_mode
         res[CONF_TRAFFIC_MODE] = self._here_data.traffic_mode
         return res
 
     @property
-    def unit_of_measurement(self):
+    def unit_of_measurement(self) -> str:
         """Return the unit this state is expressed in."""
         return self._unit_of_measurement
 
     @property
-    def icon(self):
+    def icon(self) -> str:
         """Icon to use in the frontend depending on travel_mode."""
         if self._here_data.travel_mode == TRAVEL_MODE_PEDESTRIAN:
             return ICON_PEDESTRIAN
@@ -201,18 +210,16 @@ def icon(self):
             return ICON_TRUCK
         return ICON_CAR
 
-    async def async_update(self):
+    async def async_update(self) -> None:
         """Update Sensor Information."""
         # Convert device_trackers to HERE friendly location
         if self._origin_entity_id is not None:
             self._here_data.origin = await self._get_location_from_entity(
-                self._origin_entity_id
-            )
+                self._origin_entity_id)
 
         if self._destination_entity_id is not None:
             self._here_data.destination = await self._get_location_from_entity(
-                self._destination_entity_id
-            )
+                self._destination_entity_id)
 
         self._here_data.destination = await self._resolve_zone(
             self._here_data.destination)
@@ -221,7 +228,7 @@ async def async_update(self):
 
         await self._hass.async_add_executor_job(self._here_data.update)
 
-    async def _get_location_from_entity(self, entity_id):
+    async def _get_location_from_entity(self, entity_id: str) -> Optional[str]:
         """Get the location from the entity state or attributes."""
         entity = self._hass.states.get(entity_id)
 
@@ -247,18 +254,16 @@ async def _get_location_from_entity(self, entity_id):
         if entity_id.startswith("sensor."):
             return entity.state
 
-        # When everything fails just return nothing
-        return None
-
     @staticmethod
-    def _get_location_from_attributes(entity):
+    def _get_location_from_attributes(entity: State) -> str:
         """Get the lat/long string from an entities attributes."""
         attr = entity.attributes
         return "{},{}".format(
             attr.get(ATTR_LATITUDE), attr.get(ATTR_LONGITUDE)
         )
 
-    async def _resolve_zone(self, friendly_name):
+    async def _resolve_zone(self, friendly_name: str) -> str:
+        """Get the lat/long string of a zone given its friendly_name."""
         entities = self._hass.states.async_all()
         for entity in entities:
             if entity.domain == 'zone' and entity.name == friendly_name:
@@ -270,8 +275,16 @@ async def _resolve_zone(self, friendly_name):
 class HERETravelTimeData():
     """HERETravelTime data object."""
 
-    def __init__(self, origin, destination, app_id, app_code, travel_mode,
-                 traffic_mode, route_mode, units):
+    def __init__(
+            self,
+            origin: None,
+            destination: None,
+            app_id: str,
+            app_code: str,
+            travel_mode: str,
+            traffic_mode: bool,
+            route_mode: str,
+            units: str) -> None:
         """Initialize herepy."""
         self.origin = origin
         self.destination = destination
@@ -319,8 +332,8 @@ def update(self):
                 [self.travel_mode, self.route_mode, traffic_mode],
             )
             if isinstance(response, herepy.error.HEREError):
-                # Better error message for cryptic error code
-                if 'NGEO_ERROR_GRAPH_DISCONNECTED' in response.message:
+                # Better error message for cryptic no route error codes
+                if any(error in response.message for error in NO_ROUTE_ERRORS):
                     _LOGGER.error(NO_ROUTE_ERROR_MESSAGE)
                 else:
                     _LOGGER.error("API returned error %s", response.message)