Merge pull request #23 from dev-stb/feature/dataset-filter

feat: ✨ Add options to filter the dataset.

custom_components/deutschebahn/config_flow.py

@@ -1,24 +1,71 @@
 """Config flow"""
 import logging
+from typing import Any
 
 import voluptuous as vol
 
 from homeassistant import config_entries
-from homeassistant.const import CONF_NAME
 from homeassistant.core import callback
+from homeassistant.data_entry_flow import FlowResult
 import homeassistant.helpers.config_validation as cv
 
 from .const import (  # pylint: disable=unused-import
     CONF_DESTINATION,
     CONF_START,
     CONF_OFFSET,
     CONF_ONLY_DIRECT,
+    CONF_MAX_CONNECTIONS,
+    CONF_IGNORED_PRODUCTS,
+    CONF_IGNORED_PRODUCTS_OPTIONS,
 )
+
 DOMAIN = "deutschebahn"
 
 _LOGGER = logging.getLogger(__name__)
 
 
+class OptionsFlowHandler(config_entries.OptionsFlow):
+    def __init__(self, config_entry: config_entries.ConfigEntry) -> None:
+        """Initialize options flow."""
+        self.config_entry = config_entry
+
+    async def async_step_init(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Manage the options."""
+
+        def __get_option(key: str, default: Any) -> Any:
+            return self.config_entry.options.get(
+                key, self.config_entry.data.get(key, default)
+            )
+
+        if user_input is not None:
+            return self.async_create_entry(data=user_input)
+
+        return self.async_show_form(
+            step_id="init",
+            data_schema=vol.Schema(
+                {
+                    vol.Required(
+                        CONF_OFFSET, default=__get_option(CONF_OFFSET, 0)
+                    ): cv.positive_int,
+                    vol.Required(
+                        CONF_MAX_CONNECTIONS,
+                        default=__get_option(CONF_MAX_CONNECTIONS, 2),
+                    ): cv.positive_int,
+                    vol.Required(
+                        CONF_IGNORED_PRODUCTS,
+                        default=__get_option(CONF_IGNORED_PRODUCTS, []),
+                    ): cv.multi_select(CONF_IGNORED_PRODUCTS_OPTIONS),
+                    vol.Required(
+                        CONF_ONLY_DIRECT,
+                        default=__get_option(CONF_ONLY_DIRECT, False),
+                    ): cv.boolean,
+                }
+            ),
+        )
+
+
 class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     """Handle a config flow"""
 
@@ -40,13 +87,28 @@ async def async_step_user(self, user_input=None):
 
         data_schema = vol.Schema(
             {
-                vol.Required(CONF_START): str,
-                vol.Required(CONF_DESTINATION): str,
+                vol.Required(CONF_START): cv.string,
+                vol.Required(CONF_DESTINATION): cv.string,
                 vol.Required(CONF_OFFSET, default=0): cv.positive_int,
+                vol.Required(CONF_MAX_CONNECTIONS, default=2): cv.positive_int,
+                vol.Required(
+                    CONF_IGNORED_PRODUCTS,
+                    default=[],
+                ): cv.multi_select(CONF_IGNORED_PRODUCTS_OPTIONS),
                 vol.Required(CONF_ONLY_DIRECT, default=False): cv.boolean,
-            },
+            }
         )
 
         return self.async_show_form(
-            step_id="user", data_schema=data_schema, errors=errors
-        )
+            step_id="user",
+            data_schema=data_schema,
+            errors=errors,
+        )
+
+    @staticmethod
+    @callback
+    def async_get_options_flow(
+        config_entry: config_entries.ConfigEntry,
+    ) -> config_entries.OptionsFlow:
+        """Create the options flow."""
+        return OptionsFlowHandler(config_entry)

custom_components/deutschebahn/const.py

@@ -6,5 +6,16 @@
 CONF_START = "start"
 CONF_OFFSET = "offset"
 CONF_ONLY_DIRECT = "only_direct"
-
-ATTR_DATA = "data"
+CONF_MAX_CONNECTIONS = "max_connections"
+CONF_IGNORED_PRODUCTS = "ignored_products"
+CONF_IGNORED_PRODUCTS_OPTIONS = {
+    "STR": "Straßenbahn (STR)",
+    "S": "Stadtbahn (S-Bahn)",
+    "RE": "Regional Express (RE)",
+    "RB": "Regional Bahn (RB)",
+    "EC": "EuroCity (EC)",
+    "IC": "Intercity (IC)",
+    "ICE": "Intercity Express (ICE)",
+    "TVG": "Train à grande vitesse (TVG)",
+}
+ATTR_DATA = "data"

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,8 +26,9 @@
     CONF_START,
     CONF_OFFSET,
     CONF_ONLY_DIRECT,
+    CONF_MAX_CONNECTIONS,
+    CONF_IGNORED_PRODUCTS,
     ATTR_DATA,
-
     DOMAIN,
 )
 
@@ -62,6 +64,8 @@ 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.get(CONF_MAX_CONNECTIONS, 2)
+        self.ignored_products = config.get(CONF_IGNORED_PRODUCTS, [])
         self.only_direct = config[CONF_ONLY_DIRECT]
         self.schiene = schiene.Schiene()
         self.connections = [{}]
@@ -98,7 +102,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 +113,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 +158,22 @@ 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
+        elif set(connection["products"]).intersection(self.ignored_products):
+            continue
+
+        data.append(connection)
+    _LOGGER.debug(f"Filtered data: {data}")
+
     return data

custom_components/deutschebahn/strings.json

@@ -1,4 +1,18 @@
 {
+  "options": {
+    "flow_title": "Setup DB:{entity_name}",
+    "step": {
+      "init": {
+        "description": "Submit your train station details for the search queries.",
+        "data": {
+          "offset": "Offset in minutes",
+          "only_direct": "Only show direct connections?",
+          "max_connections": "Max connections in sensor",
+          "ignored_products": "Products to ignore"
+        }
+      }
+    }
+  },
   "config": {
     "flow_title": "Setup DeutscheBahn",
     "step": {
@@ -8,7 +22,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",
+          "ignored_products": "Products to ignore"
         }
       }
     },

custom_components/deutschebahn/translations/de.json

@@ -1,4 +1,18 @@
 {
+  "options": {
+    "flow_title": "Einrichtung der:{entity_name}",
+    "step": {
+      "init": {
+        "description": "Trage hier deine Zug Station Informationen ein für die Suchabfrage.",
+        "data": {
+          "offset": "Versatz in Minutes",
+          "only_direct": "Zeige nur direkte Verbindungen?",
+          "max_connections": "Maximale Anzahle der angezeigten Verbindungen",
+          "ignored_products": "Zu ignorierende Produkte"
+        }
+      }
+    }
+  },
   "config": {
     "flow_title": "Einrichtung der DeutscheBahn Integration",
     "step": {
@@ -8,7 +22,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",
+          "ignored_products": "Zu ignorierende Produkte"
         }
       }
     },

custom_components/deutschebahn/translations/en.json

@@ -1,4 +1,18 @@
 {
+  "options": {
+    "flow_title": "Setup DB:{entity_name}",
+    "step": {
+      "init": {
+        "description": "Submit your train station details for the search queries.",
+        "data": {
+          "offset": "Offset in minutes",
+          "only_direct": "Only show direct connections?",
+          "max_connections": "Max connections in sensor",
+          "ignored_products": "Products to ignore"
+        }
+      }
+    }
+  },
   "config": {
     "flow_title": "Setup DeutscheBahn",
     "step": {
@@ -8,7 +22,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",
+          "ignored_products": "Products to ignore"
         }
       }
     },