diff --git a/homeassistant/components/deutsche_bahn/sensor.py b/homeassistant/components/deutsche_bahn/sensor.py index 9c7518eb8ef..c6761c58e57 100644 --- a/homeassistant/components/deutsche_bahn/sensor.py +++ b/homeassistant/components/deutsche_bahn/sensor.py @@ -13,6 +13,8 @@ _LOGGER = logging.getLogger(__name__) CONF_DESTINATION = 'to' CONF_START = 'from' +CONF_OFFSET = 'offset' +DEFAULT_OFFSET = timedelta(minutes=0) CONF_ONLY_DIRECT = 'only_direct' DEFAULT_ONLY_DIRECT = False @@ -23,6 +25,7 @@ SCAN_INTERVAL = timedelta(minutes=2) PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({ vol.Required(CONF_DESTINATION): cv.string, vol.Required(CONF_START): cv.string, + vol.Optional(CONF_OFFSET, default=DEFAULT_OFFSET): cv.time_period, vol.Optional(CONF_ONLY_DIRECT, default=DEFAULT_ONLY_DIRECT): cv.boolean, }) @@ -31,18 +34,20 @@ def setup_platform(hass, config, add_entities, discovery_info=None): """Set up the Deutsche Bahn Sensor.""" start = config.get(CONF_START) destination = config.get(CONF_DESTINATION) + offset = config.get(CONF_OFFSET) only_direct = config.get(CONF_ONLY_DIRECT) - add_entities([DeutscheBahnSensor(start, destination, only_direct)], True) + add_entities([DeutscheBahnSensor(start, destination, + offset, only_direct)], True) class DeutscheBahnSensor(Entity): """Implementation of a Deutsche Bahn sensor.""" - def __init__(self, start, goal, only_direct): + def __init__(self, start, goal, offset, only_direct): """Initialize the sensor.""" self._name = '{} to {}'.format(start, goal) - self.data = SchieneData(start, goal, only_direct) + self.data = SchieneData(start, goal, offset, only_direct) self._state = None @property @@ -81,12 +86,13 @@ class DeutscheBahnSensor(Entity): class SchieneData: """Pull data from the bahn.de web page.""" - def __init__(self, start, goal, only_direct): + def __init__(self, start, goal, offset, only_direct): """Initialize the sensor.""" import schiene self.start = start self.goal = goal + self.offset = offset self.only_direct = only_direct self.schiene = schiene.Schiene() self.connections = [{}] @@ -94,7 +100,8 @@ class SchieneData: def update(self): """Update the connection data.""" self.connections = self.schiene.connections( - self.start, self.goal, dt_util.as_local(dt_util.utcnow()), + self.start, self.goal, + dt_util.as_local(dt_util.utcnow()+self.offset), self.only_direct) if not self.connections: