Add next starship launch sensor to launch_library (#64929)

Co-authored-by: Joakim Sørensen <hi@ludeeus.dev>

homeassistant/components/launch_library/__init__.py

@@ -1,8 +1,13 @@
 """The launch_library component."""
+from __future__ import annotations
+
 from datetime import timedelta
 import logging
+from typing import TypedDict
 
 from pylaunches import PyLaunches, PyLaunchesException
+from pylaunches.objects.launch import Launch
+from pylaunches.objects.starship import StarshipResponse
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
@@ -17,6 +22,13 @@ _LOGGER = logging.getLogger(__name__)
 PLATFORMS = [Platform.SENSOR]
 
 
+class LaunchLibraryData(TypedDict):
+    """Typed dict representation of data returned from pylaunches."""
+
+    upcoming_launches: list[Launch]
+    starship_events: StarshipResponse
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up this integration using UI."""
 
@@ -25,9 +37,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     session = async_get_clientsession(hass)
     launches = PyLaunches(session)
 
-    async def async_update():
+    async def async_update() -> LaunchLibraryData:
         try:
-            return await launches.upcoming_launches()
+            return LaunchLibraryData(
+                upcoming_launches=await launches.upcoming_launches(),
+                starship_events=await launches.starship_events(),
+            )
         except PyLaunchesException as ex:
             raise UpdateFailed(ex) from ex
 

homeassistant/components/launch_library/diagnostics.py

@@ -3,12 +3,11 @@ from __future__ import annotations
 
 from typing import Any
 
-from pylaunches.objects.launch import Launch
-
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
+from . import LaunchLibraryData
 from .const import DOMAIN
 
 
@@ -17,8 +16,12 @@ async def async_get_config_entry_diagnostics(
     entry: ConfigEntry,
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinator: DataUpdateCoordinator[list[Launch]] = hass.data[DOMAIN]
+    coordinator: DataUpdateCoordinator[LaunchLibraryData] = hass.data[DOMAIN]
-    next_launch = coordinator.data[0] if coordinator.data else None
+    if coordinator.data is None:
+        return {}
+    next_launch = coordinator.data["upcoming_launches"][0]
+    starship_launch = coordinator.data["starship_events"].upcoming.launches[0]
     return {
-        "next_launch": next_launch.raw_data_contents if next_launch else None,
+        "next_launch": next_launch.raw_data_contents,
+        "starship": starship_launch.raw_data_contents,
     }

homeassistant/components/launch_library/sensor.py

@@ -28,6 +28,7 @@ from homeassistant.helpers.update_coordinator import (
 )
 from homeassistant.util.dt import parse_datetime
 
+from . import LaunchLibraryData
 from .const import DOMAIN
 
 DEFAULT_NEXT_LAUNCH_NAME = "Next launch"
@@ -41,7 +42,7 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
 
 
 @dataclass
-class NextLaunchSensorEntityDescriptionMixin:
+class LaunchLibrarySensorEntityDescriptionMixin:
     """Mixin for required keys."""
 
     value_fn: Callable[[Launch], datetime | int | str | None]
@@ -49,14 +50,14 @@ class NextLaunchSensorEntityDescriptionMixin:
 
 
 @dataclass
-class NextLaunchSensorEntityDescription(
+class LaunchLibrarySensorEntityDescription(
-    SensorEntityDescription, NextLaunchSensorEntityDescriptionMixin
+    SensorEntityDescription, LaunchLibrarySensorEntityDescriptionMixin
 ):
     """Describes a Next Launch sensor entity."""
 
 
-SENSOR_DESCRIPTIONS: tuple[NextLaunchSensorEntityDescription, ...] = (
+SENSOR_DESCRIPTIONS: tuple[LaunchLibrarySensorEntityDescription, ...] = (
-    NextLaunchSensorEntityDescription(
+    LaunchLibrarySensorEntityDescription(
         key="next_launch",
         icon="mdi:rocket-launch",
         name="Next launch",
@@ -68,7 +69,7 @@ SENSOR_DESCRIPTIONS: tuple[NextLaunchSensorEntityDescription, ...] = (
             "provider_country_code": nl.pad.location.country_code,
         },
     ),
-    NextLaunchSensorEntityDescription(
+    LaunchLibrarySensorEntityDescription(
         key="launch_time",
         icon="mdi:clock-outline",
         name="Launch time",
@@ -80,7 +81,7 @@ SENSOR_DESCRIPTIONS: tuple[NextLaunchSensorEntityDescription, ...] = (
             "stream_live": nl.webcast_live,
         },
     ),
-    NextLaunchSensorEntityDescription(
+    LaunchLibrarySensorEntityDescription(
         key="launch_probability",
         icon="mdi:dice-multiple",
         name="Launch Probability",
@@ -88,14 +89,14 @@ SENSOR_DESCRIPTIONS: tuple[NextLaunchSensorEntityDescription, ...] = (
         value_fn=lambda nl: None if nl.probability == -1 else nl.probability,
         attributes_fn=lambda nl: None,
     ),
-    NextLaunchSensorEntityDescription(
+    LaunchLibrarySensorEntityDescription(
         key="launch_status",
         icon="mdi:rocket-launch",
         name="Launch status",
         value_fn=lambda nl: nl.status.name,
         attributes_fn=lambda nl: {"reason": nl.holdreason} if nl.inhold else None,
     ),
-    NextLaunchSensorEntityDescription(
+    LaunchLibrarySensorEntityDescription(
         key="launch_mission",
         icon="mdi:orbit",
         name="Launch mission",
@@ -106,6 +107,19 @@ SENSOR_DESCRIPTIONS: tuple[NextLaunchSensorEntityDescription, ...] = (
             "description": nl.mission.description,
         },
     ),
+    LaunchLibrarySensorEntityDescription(
+        key="starship_launch",
+        icon="mdi:rocket",
+        name="Next Starship launch",
+        device_class=SensorDeviceClass.TIMESTAMP,
+        value_fn=lambda sl: parse_datetime(sl.net),
+        attributes_fn=lambda sl: {
+            "title": sl.mission.name,
+            "status": sl.status.name,
+            "target_orbit": sl.mission.orbit.name,
+            "description": sl.mission.description,
+        },
+    ),
 )
 
 
@@ -138,10 +152,10 @@ async def async_setup_entry(
 ) -> None:
     """Set up the sensor platform."""
     name = entry.data.get(CONF_NAME, DEFAULT_NEXT_LAUNCH_NAME)
-    coordinator = hass.data[DOMAIN]
+    coordinator: DataUpdateCoordinator[LaunchLibraryData] = hass.data[DOMAIN]
 
     async_add_entities(
-        NextLaunchSensor(
+        LaunchLibrarySensor(
             coordinator=coordinator,
             entry_id=entry.entry_id,
             description=description,
@@ -151,18 +165,19 @@ async def async_setup_entry(
     )
 
 
-class NextLaunchSensor(CoordinatorEntity, SensorEntity):
+class LaunchLibrarySensor(CoordinatorEntity, SensorEntity):
     """Representation of the next launch sensors."""
 
     _attr_attribution = "Data provided by Launch Library."
     _next_launch: Launch | None = None
-    entity_description: NextLaunchSensorEntityDescription
+    entity_description: LaunchLibrarySensorEntityDescription
+    coordinator: DataUpdateCoordinator[LaunchLibraryData]
 
     def __init__(
         self,
-        coordinator: DataUpdateCoordinator,
+        coordinator: DataUpdateCoordinator[LaunchLibraryData],
         entry_id: str,
-        description: NextLaunchSensorEntityDescription,
+        description: LaunchLibrarySensorEntityDescription,
         name: str | None = None,
     ) -> None:
         """Initialize a Launch Library sensor."""
@@ -194,7 +209,12 @@ class NextLaunchSensor(CoordinatorEntity, SensorEntity):
     @callback
     def _handle_coordinator_update(self) -> None:
         """Handle updated data from the coordinator."""
-        self._next_launch = next((launch for launch in self.coordinator.data), None)
+        if self.entity_description.key == "starship_launch":
+            launches = self.coordinator.data["starship_events"].upcoming.launches
+        else:
+            launches = self.coordinator.data["upcoming_launches"]
+
+        self._next_launch = next((launch for launch in (launches)), None)
         super()._handle_coordinator_update()
 
     async def async_added_to_hass(self) -> None: