add service to fetch more connections

homeassistant/components/swiss_public_transport/const.py

@@ -1,5 +1,7 @@
 """Constants for the swiss_public_transport integration."""
 
+from typing import Final
+
 DOMAIN = "swiss_public_transport"
 
 CONF_DESTINATION = "to"
@@ -8,9 +10,14 @@
 DEFAULT_NAME = "Next Destination"
 
 SENSOR_CONNECTIONS_COUNT = 3
+SENSOR_CONNECTIONS_MAX = 15
 
 
 PLACEHOLDERS = {
     "stationboard_url": "http://transport.opendata.ch/examples/stationboard.html",
     "opendata_url": "http://transport.opendata.ch",
 }
+
+ATTR_LIMIT: Final = "limit"
+
+SERVICE_FETCH_CONNECTIONS = "fetch_connections"

homeassistant/components/swiss_public_transport/coordinator.py

@@ -76,6 +76,11 @@ def nth_departure_time(self, i: int) -> datetime | None:
         return None
 
     async def _async_update_data(self) -> list[DataConnection]:
+        return await self.fetch_connections(limit=SENSOR_CONNECTIONS_COUNT)
+
+    async def fetch_connections(self, limit: int) -> list[DataConnection]:
+        """Fetch connections using the opendata api."""
+        self._opendata.limit = limit
         try:
             await self._opendata.async_get_data()
         except OpendataTransportError as e:
@@ -99,6 +104,6 @@ async def _async_update_data(self) -> list[DataConnection]:
                 remaining_time=str(self.remaining_time(connections[i]["departure"])),
                 delay=connections[i]["delay"],
             )
-            for i in range(SENSOR_CONNECTIONS_COUNT)
+            for i in range(limit)
             if len(connections) > i and connections[i] is not None
         ]

homeassistant/components/swiss_public_transport/icons.json

@@ -23,5 +23,8 @@
         "default": "mdi:clock-plus"
       }
     }
+  },
+  "services": {
+    "fetch_connections": "mdi:bus-clock"
   }
 }

homeassistant/components/swiss_public_transport/sensor.py

@@ -8,33 +8,50 @@
 import logging
 from typing import TYPE_CHECKING
 
+from opendata_transport.exceptions import OpendataTransportError
 import voluptuous as vol
 
 from homeassistant import config_entries, core
 from homeassistant.components.sensor import (
+    DOMAIN as SENSOR_DOMAIN,
     PLATFORM_SCHEMA,
     SensorDeviceClass,
     SensorEntity,
     SensorEntityDescription,
 )
 from homeassistant.config_entries import SOURCE_IMPORT
-from homeassistant.const import CONF_NAME, UnitOfTime
-from homeassistant.core import DOMAIN as HOMEASSISTANT_DOMAIN, HomeAssistant, callback
+from homeassistant.const import ATTR_ENTITY_ID, CONF_NAME, UnitOfTime
+from homeassistant.core import (
+    DOMAIN as HOMEASSISTANT_DOMAIN,
+    HomeAssistant,
+    SupportsResponse,
+    callback,
+)
 from homeassistant.data_entry_flow import FlowResultType
+from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers import entity_platform
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.issue_registry import IssueSeverity, async_create_issue
+from homeassistant.helpers.selector import (
+    NumberSelector,
+    NumberSelectorConfig,
+    NumberSelectorMode,
+)
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType, StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import (
+    ATTR_LIMIT,
     CONF_DESTINATION,
     CONF_START,
     DEFAULT_NAME,
     DOMAIN,
     PLACEHOLDERS,
     SENSOR_CONNECTIONS_COUNT,
+    SENSOR_CONNECTIONS_MAX,
+    SERVICE_FETCH_CONNECTIONS,
 )
 from .coordinator import DataConnection, SwissPublicTransportDataUpdateCoordinator
 
@@ -117,6 +134,26 @@ async def async_setup_entry(
         for description in SENSORS
     )
 
+    platform = entity_platform.async_get_current_platform()
+
+    platform.async_register_entity_service(
+        SERVICE_FETCH_CONNECTIONS,
+        vol.Schema(
+            {
+                vol.Required(ATTR_ENTITY_ID): cv.entities_domain(SENSOR_DOMAIN),
+                vol.Required(
+                    ATTR_LIMIT, default=SENSOR_CONNECTIONS_COUNT
+                ): NumberSelector(
+                    NumberSelectorConfig(
+                        min=0, max=SENSOR_CONNECTIONS_MAX, mode=NumberSelectorMode.BOX
+                    )
+                ),
+            }
+        ),
+        "async_fetch_connections",
+        supports_response=SupportsResponse.ONLY,
+    )
+
 
 async def async_setup_platform(
     hass: HomeAssistant,
@@ -218,3 +255,16 @@ def _async_update_attrs(self) -> None:
                 ].items()
                 if key not in {"departure"}
             }
+
+    async def async_fetch_connections(
+        self,
+        limit: int,
+    ) -> dict:
+        """Fetch a set of connections."""
+        try:
+            connections = await self.coordinator.fetch_connections(limit=int(limit))
+        except OpendataTransportError as e:
+            raise HomeAssistantError(
+                "Unable to fetch connections for swiss public trainsport"
+            ) from e
+        return {"connections": connections}

homeassistant/components/swiss_public_transport/services.yaml

@@ -0,0 +1,17 @@
+fetch_connections:
+  fields:
+    entity_id:
+      required: true
+      example: sensor.lausanne_zurich
+      selector:
+        entity:
+          domain: sensor
+          integration: swiss_public_transport
+    limit:
+      example: 3
+      required: true
+      selector:
+        number:
+          min: 0
+          max: 15
+          step: 1

homeassistant/components/swiss_public_transport/strings.json

@@ -47,6 +47,22 @@
       }
     }
   },
+  "services": {
+    "fetch_connections": {
+      "name": "Fetch Connections",
+      "description": "Fetch a list of connections from the swiss public transport api.",
+      "fields": {
+        "entity_id": {
+          "name": "List",
+          "description": "Swiss public transport sensor to fetch connections for."
+        },
+        "limit": {
+          "name": "Limit",
+          "description": "Number of connections to fetch from [1-15]"
+        }
+      }
+    }
+  },
   "issues": {
     "deprecated_yaml_import_issue_cannot_connect": {
       "title": "The swiss public transport YAML configuration import cannot connect to server",