From cb3c9512068790be486fe79b93bf21a2ac647808 Mon Sep 17 00:00:00 2001 From: Simon Bonzelet Date: Sun, 5 Nov 2023 12:08:58 +0100 Subject: [PATCH] feat: :sparkles: Add options to filter the dataset. - Allow connections list to be filtered by product - Set a maximum connection count to be available to a sensor - Allow to access connections list in sensor resolves FaserF/ha-deutschebahn#12 --- custom_components/deutschebahn/config_flow.py | 4 ++ custom_components/deutschebahn/const.py | 4 +- custom_components/deutschebahn/sensor.py | 39 ++++++++++++++++--- custom_components/deutschebahn/strings.json | 4 +- .../deutschebahn/translations/de.json | 4 +- .../deutschebahn/translations/en.json | 4 +- 6 files changed, 50 insertions(+), 9 deletions(-) diff --git a/custom_components/deutschebahn/config_flow.py b/custom_components/deutschebahn/config_flow.py index 6a005f2..e96d9ac 100644 --- a/custom_components/deutschebahn/config_flow.py +++ b/custom_components/deutschebahn/config_flow.py @@ -13,6 +13,8 @@ CONF_START, CONF_OFFSET, CONF_ONLY_DIRECT, + CONF_MAX_CONNECTIONS, + CONF_SELECTED_PRODUCTS, ) DOMAIN = "deutschebahn" @@ -43,6 +45,8 @@ async def async_step_user(self, user_input=None): vol.Required(CONF_START): str, vol.Required(CONF_DESTINATION): str, vol.Required(CONF_OFFSET, default=0): cv.positive_int, + vol.Required(CONF_MAX_CONNECTIONS, default=2): cv.positive_int, + vol.Required(CONF_SELECTED_PRODUCTS, default="All"): str, vol.Required(CONF_ONLY_DIRECT, default=False): cv.boolean, }, ) diff --git a/custom_components/deutschebahn/const.py b/custom_components/deutschebahn/const.py index 466d7d1..f5e379e 100644 --- a/custom_components/deutschebahn/const.py +++ b/custom_components/deutschebahn/const.py @@ -6,5 +6,7 @@ CONF_START = "start" CONF_OFFSET = "offset" CONF_ONLY_DIRECT = "only_direct" +CONF_MAX_CONNECTIONS = "max_connections" +CONF_SELECTED_PRODUCTS = "selected_products" -ATTR_DATA = "data" \ No newline at end of file +ATTR_DATA = "data" diff --git a/custom_components/deutschebahn/sensor.py b/custom_components/deutschebahn/sensor.py index 2c11a19..5800a6a 100644 --- a/custom_components/deutschebahn/sensor.py +++ b/custom_components/deutschebahn/sensor.py @@ -1,7 +1,8 @@ """deutschebahn sensor platform.""" from datetime import timedelta, datetime import logging -from typing import Any, Callable, Dict, Optional +from typing import Any, Callable, Dict, Optional, Set +import re import schiene import async_timeout @@ -25,6 +26,8 @@ CONF_START, CONF_OFFSET, CONF_ONLY_DIRECT, + CONF_MAX_CONNECTIONS, + CONF_SELECTED_PRODUCTS, ATTR_DATA, DOMAIN, @@ -62,6 +65,20 @@ def __init__(self, config, hass: HomeAssistantType): self.start = config[CONF_START] self.goal = config[CONF_DESTINATION] self.offset = timedelta(minutes=config[CONF_OFFSET]) + self.max_connections: int = config[CONF_MAX_CONNECTIONS] + self.select_any_product: bool = any( + [ + p.strip().upper() in ["*", "ALL", "ANY"] + for p in config[CONF_SELECTED_PRODUCTS] + ] + ) + self.selected_products: Set[str] = set( + [ + product.strip() + for product in re.split(r", |,|\s", config[CONF_SELECTED_PRODUCTS]) + if not product.isspace() + ] + ) self.only_direct = config[CONF_ONLY_DIRECT] self.schiene = schiene.Schiene() self.connections = [{}] @@ -98,7 +115,7 @@ def native_value(self): def extra_state_attributes(self): """Return the state attributes.""" if len(self.connections) > 0: - connections = self.connections[0] + connections = self.connections[0].copy() if len(self.connections) > 1: connections["next"] = self.connections[1]["departure"] connections["next_delay"] = self.connections[1]["delay"] @@ -109,6 +126,7 @@ def extra_state_attributes(self): connections["next_on_canceled"] = self.connections[2]["canceled"] else: connections = None + connections["departures"] = self.connections return connections async def async_update(self): @@ -153,12 +171,23 @@ async def async_update(self): _LOGGER.exception(f"Cannot retrieve data for direction: '{self.start}' '{self.goal}'") def fetch_schiene_connections(hass, self): - _LOGGER.debug(f"Fetching update from schiene python module for '{self.start}' '{self.goal}'") - data = self.schiene.connections( + raw_data = self.schiene.connections( self.start, self.goal, dt_util.as_local(dt_util.utcnow() + self.offset), self.only_direct, ) - _LOGGER.debug(f"Fetched data: {data}") + _LOGGER.debug(f"Fetched data: {raw_data}") + + data = [] + for connection in raw_data: + if len(data) == self.max_connections: + break + if self.select_any_product or set(connection["products"]).intersection( + self.selected_products + ): + data.append(connection) + + _LOGGER.debug(f"Filtered data: {data}") + return data diff --git a/custom_components/deutschebahn/strings.json b/custom_components/deutschebahn/strings.json index 8226e15..a929ab5 100644 --- a/custom_components/deutschebahn/strings.json +++ b/custom_components/deutschebahn/strings.json @@ -8,7 +8,9 @@ "start": "Start station", "destination": "Destination station", "offset": "Offset in minutes", - "only_direct": "Only show direct connections?" + "only_direct": "Only show direct connections?", + "max_connections": "Max connections in sensor", + "selected_products": "Product selection (comma seperated)" } } }, diff --git a/custom_components/deutschebahn/translations/de.json b/custom_components/deutschebahn/translations/de.json index b33a7f1..fe8d919 100644 --- a/custom_components/deutschebahn/translations/de.json +++ b/custom_components/deutschebahn/translations/de.json @@ -8,7 +8,9 @@ "start": "Start Station", "destination": "Ziel station", "offset": "Versatz in Minutes", - "only_direct": "Zeige nur direkte Verbindungen?" + "only_direct": "Zeige nur direkte Verbindungen?", + "max_connections": "Maximale Anzahle der angezeigten Verbindungen", + "selected_products": "Produktauswahl" } } }, diff --git a/custom_components/deutschebahn/translations/en.json b/custom_components/deutschebahn/translations/en.json index 8226e15..a929ab5 100644 --- a/custom_components/deutschebahn/translations/en.json +++ b/custom_components/deutschebahn/translations/en.json @@ -8,7 +8,9 @@ "start": "Start station", "destination": "Destination station", "offset": "Offset in minutes", - "only_direct": "Only show direct connections?" + "only_direct": "Only show direct connections?", + "max_connections": "Max connections in sensor", + "selected_products": "Product selection (comma seperated)" } } },