From 4e808133f224053cb9f66b1b2b24dc09264cbd2c Mon Sep 17 00:00:00 2001 From: Simon Hansen <67142049+DurgNomis-drol@users.noreply.github.com> Date: Wed, 26 Jan 2022 12:31:03 +0100 Subject: [PATCH] Add next starship launch sensor to launch_library (#64929) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Joakim Sørensen --- .../components/launch_library/__init__.py | 19 ++++++- .../components/launch_library/diagnostics.py | 13 +++-- .../components/launch_library/sensor.py | 52 +++++++++++++------ 3 files changed, 61 insertions(+), 23 deletions(-) diff --git a/homeassistant/components/launch_library/__init__.py b/homeassistant/components/launch_library/__init__.py index db4512b0885..14eac4fd6a6 100644 --- a/homeassistant/components/launch_library/__init__.py +++ b/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 diff --git a/homeassistant/components/launch_library/diagnostics.py b/homeassistant/components/launch_library/diagnostics.py index 18744e49b39..362ab561a31 100644 --- a/homeassistant/components/launch_library/diagnostics.py +++ b/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] - next_launch = coordinator.data[0] if coordinator.data else None + coordinator: DataUpdateCoordinator[LaunchLibraryData] = hass.data[DOMAIN] + 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, } diff --git a/homeassistant/components/launch_library/sensor.py b/homeassistant/components/launch_library/sensor.py index 668e6a9b75a..61c19a0dfb6 100644 --- a/homeassistant/components/launch_library/sensor.py +++ b/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( - SensorEntityDescription, NextLaunchSensorEntityDescriptionMixin +class LaunchLibrarySensorEntityDescription( + SensorEntityDescription, LaunchLibrarySensorEntityDescriptionMixin ): """Describes a Next Launch sensor entity.""" -SENSOR_DESCRIPTIONS: tuple[NextLaunchSensorEntityDescription, ...] = ( - NextLaunchSensorEntityDescription( +SENSOR_DESCRIPTIONS: tuple[LaunchLibrarySensorEntityDescription, ...] = ( + 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: