Add the two next trains (#2390)

homeassistant/components/sensor/deutsche_bahn.py

@@ -1,38 +1,43 @@
 """
-Support for information about the German trans system.
+Support for information about the German train system.
 
 For more details about this platform, please refer to the documentation at
 https://home-assistant.io/components/sensor.deutsche_bahn/
 """
 import logging
-from datetime import timedelta, datetime
+from datetime import timedelta
+
+import voluptuous as vol
+
+from homeassistant.const import (CONF_PLATFORM)
+import homeassistant.helpers.config_validation as cv
 from homeassistant.util import Throttle
 from homeassistant.helpers.entity import Entity
 
-_LOGGER = logging.getLogger(__name__)
 REQUIREMENTS = ['schiene==0.17']
+
+CONF_START = 'from'
+CONF_DESTINATION = 'to'
 ICON = 'mdi:train'
 
+_LOGGER = logging.getLogger(__name__)
+
+PLATFORM_SCHEMA = vol.Schema({
+    vol.Required(CONF_PLATFORM): 'deutsche_bahn',
+    vol.Required(CONF_START): cv.string,
+    vol.Required(CONF_DESTINATION): cv.string,
+})
+
 # Return cached results if last scan was less then this time ago.
 MIN_TIME_BETWEEN_UPDATES = timedelta(seconds=120)
 
 
-def setup_platform(hass, config, add_devices_callback, discovery_info=None):
+def setup_platform(hass, config, add_devices, discovery_info=None):
     """Setup the Deutsche Bahn Sensor."""
-    start = config.get('from')
-    goal = config.get('to')
+    start = config.get(CONF_START)
+    destination = config.get(CONF_DESTINATION)
 
-    if start is None:
-        _LOGGER.error('Missing required variable: "from"')
-        return False
-
-    if goal is None:
-        _LOGGER.error('Missing required variable: "to"')
-        return False
-
-    dev = []
-    dev.append(DeutscheBahnSensor(start, goal))
-    add_devices_callback(dev)
+    add_devices([DeutscheBahnSensor(start, destination)])
 
 
 # pylint: disable=too-few-public-methods
@@ -63,16 +68,17 @@ class DeutscheBahnSensor(Entity):
     @property
     def state_attributes(self):
         """Return the state attributes."""
-        return self.data.connections[0]
+        connections = self.data.connections[0]
+        connections['next'] = self.data.connections[1]['departure']
+        connections['next_on'] = self.data.connections[2]['departure']
+        return connections
 
     def update(self):
         """Get the latest delay from bahn.de and updates the state."""
         self.data.update()
         self._state = self.data.connections[0].get('departure', 'Unknown')
         if self.data.connections[0]['delay'] != 0:
-            self._state += " + {}".format(
-                self.data.connections[0]['delay']
-            )
+            self._state += " + {}".format(self.data.connections[0]['delay'])
 
 
 # pylint: disable=too-few-public-methods
@@ -90,18 +96,15 @@ class SchieneData(object):
     @Throttle(MIN_TIME_BETWEEN_UPDATES)
     def update(self):
         """Update the connection data."""
-        self.connections = self.schiene.connections(self.start,
-                                                    self.goal,
-                                                    datetime.now())
+        self.connections = self.schiene.connections(self.start, self.goal)
+
         for con in self.connections:
-            # Details info is not useful.
-            # Having a more consistent interface simplifies
-            # usage of Template sensors later on
+            # Detail info is not useful. Having a more consistent interface
+            # simplifies usage of template sensors.
             if 'details' in con:
                 con.pop('details')
-                delay = con.get('delay',
-                                {'delay_departure': 0,
-                                 'delay_arrival': 0})
-                # IMHO only delay_departure is usefull
+                delay = con.get('delay', {'delay_departure': 0,
+                                          'delay_arrival': 0})
+                # IMHO only delay_departure is useful
                 con['delay'] = delay['delay_departure']
                 con['ontime'] = con.get('ontime', False)