diff --git a/homeassistant/components/sensor/mvglive.py b/homeassistant/components/sensor/mvglive.py index 95f2845eced6d3..e066bb5e0b9b4f 100644 --- a/homeassistant/components/sensor/mvglive.py +++ b/homeassistant/components/sensor/mvglive.py @@ -7,6 +7,7 @@ import logging from datetime import timedelta +from copy import deepcopy import voluptuous as vol import homeassistant.helpers.config_validation as cv @@ -28,6 +29,7 @@ CONF_LINES = 'lines' CONF_PRODUCTS = 'products' CONF_TIMEOFFSET = 'timeoffset' +CONF_NUMBER = 'number' DEFAULT_PRODUCT = ['U-Bahn', 'Tram', 'Bus', 'S-Bahn'] @@ -52,6 +54,7 @@ vol.Optional(CONF_PRODUCTS, default=DEFAULT_PRODUCT): cv.ensure_list_csv, vol.Optional(CONF_TIMEOFFSET, default=0): cv.positive_int, + vol.Optional(CONF_NUMBER, default=1): cv.positive_int, vol.Optional(CONF_NAME): cv.string}] }) @@ -68,6 +71,7 @@ def setup_platform(hass, config, add_devices, discovery_info=None): nextdeparture.get(CONF_LINES), nextdeparture.get(CONF_PRODUCTS), nextdeparture.get(CONF_TIMEOFFSET), + nextdeparture.get(CONF_NUMBER), nextdeparture.get(CONF_NAME))) add_devices(sensors, True) @@ -76,12 +80,12 @@ class MVGLiveSensor(Entity): """Implementation of an MVG Live sensor.""" def __init__(self, station, destinations, directions, - lines, products, timeoffset, name): + lines, products, timeoffset, number, name): """Initialize the sensor.""" self._station = station self._name = name self.data = MVGLiveData(station, destinations, directions, - lines, products, timeoffset) + lines, products, timeoffset, number) self._state = STATE_UNKNOWN self._icon = ICONS['-'] @@ -98,9 +102,14 @@ def state(self): return self._state @property - def state_attributes(self): + def device_state_attributes(self): """Return the state attributes.""" - return self.data.departures + dep = self.data.departures + if not dep: + return None + attr = dep[0] # next depature attributes + attr['departures'] = deepcopy(dep) # all departures dictionary + return attr @property def icon(self): @@ -119,15 +128,15 @@ def update(self): self._state = '-' self._icon = ICONS['-'] else: - self._state = self.data.departures.get('time', '-') - self._icon = ICONS[self.data.departures.get('product', '-')] + self._state = self.data.departures[0].get('time', '-') + self._icon = ICONS[self.data.departures[0].get('product', '-')] class MVGLiveData: """Pull data from the mvg-live.de web page.""" def __init__(self, station, destinations, directions, - lines, products, timeoffset): + lines, products, timeoffset, number): """Initialize the sensor.""" import MVGLive self._station = station @@ -136,25 +145,30 @@ def __init__(self, station, destinations, directions, self._lines = lines self._products = products self._timeoffset = timeoffset + self._number = number self._include_ubahn = True if 'U-Bahn' in self._products else False self._include_tram = True if 'Tram' in self._products else False self._include_bus = True if 'Bus' in self._products else False self._include_sbahn = True if 'S-Bahn' in self._products else False self.mvg = MVGLive.MVGLive() - self.departures = {} + self.departures = [] def update(self): """Update the connection data.""" try: _departures = self.mvg.getlivedata( - station=self._station, ubahn=self._include_ubahn, - tram=self._include_tram, bus=self._include_bus, + station=self._station, + timeoffset=self._timeoffset, + ubahn=self._include_ubahn, + tram=self._include_tram, + bus=self._include_bus, sbahn=self._include_sbahn) except ValueError: - self.departures = {} + self.departures = [] _LOGGER.warning("Returned data not understood") return - for _departure in _departures: + self.departures = [] + for i, _departure in enumerate(_departures): # find the first departure meeting the criteria if ('' not in self._destinations[:1] and _departure['destination'] not in self._destinations): @@ -173,5 +187,6 @@ def update(self): 'product']: _nextdep[k] = _departure.get(k, '') _nextdep['time'] = int(_nextdep['time']) - self.departures = _nextdep - break + self.departures.append(_nextdep) + if i == self._number - 1: + break