Add two new sensors to Launch Library (#64203)

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

homeassistant/components/launch_library/sensor.py

@@ -1,6 +1,9 @@
 """Support for Launch Library sensors."""
 from __future__ import annotations
 
+from collections.abc import Callable
+from dataclasses import dataclass
+from datetime import datetime
 import logging
 from typing import Any
 
@@ -9,11 +12,12 @@ import voluptuous as vol
 
 from homeassistant.components.sensor import (
     PLATFORM_SCHEMA,
+    SensorDeviceClass,
     SensorEntity,
     SensorEntityDescription,
 )
 from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
-from homeassistant.const import CONF_NAME
+from homeassistant.const import CONF_NAME, PERCENTAGE, STATE_UNKNOWN
 from homeassistant.core import HomeAssistant, callback
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -22,15 +26,22 @@ from homeassistant.helpers.update_coordinator import (
     CoordinatorEntity,
     DataUpdateCoordinator,
 )
+from homeassistant.util.dt import parse_datetime
 
 from .const import (
-    ATTR_AGENCY,
-    ATTR_AGENCY_COUNTRY_CODE,
-    ATTR_LAUNCH_TIME,
-    ATTR_STREAM,
+    ATTR_LAUNCH_FACILITY,
+    ATTR_LAUNCH_PAD,
+    ATTR_LAUNCH_PAD_COUNTRY_CODE,
+    ATTR_LAUNCH_PROVIDER,
+    ATTR_STREAM_LIVE,
+    ATTR_WINDOW_END,
+    ATTR_WINDOW_START,
     ATTRIBUTION,
     DEFAULT_NAME,
     DOMAIN,
+    LAUNCH_PROBABILITY,
+    LAUNCH_TIME,
+    NEXT_LAUNCH,
 )
 
 _LOGGER = logging.getLogger(__name__)
@@ -41,10 +52,56 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
 )
 
 
-SENSOR_DESCRIPTION = SensorEntityDescription(
-    key="next_launch",
-    icon="mdi:rocket-launch",
-    name=DEFAULT_NAME,
+@dataclass
+class NextLaunchSensorEntityDescriptionMixin:
+    """Mixin for required keys."""
+
+    value_fn: Callable[[Launch], datetime | int | str | None]
+    attributes_fn: Callable[[Launch], dict[str, Any] | None]
+
+
+@dataclass
+class NextLaunchSensorEntityDescription(
+    SensorEntityDescription, NextLaunchSensorEntityDescriptionMixin
+):
+    """Describes a Next Launch sensor entity."""
+
+
+SENSOR_DESCRIPTIONS: tuple[NextLaunchSensorEntityDescription, ...] = (
+    NextLaunchSensorEntityDescription(
+        key=NEXT_LAUNCH,
+        icon="mdi:rocket-launch",
+        name=DEFAULT_NAME,
+        value_fn=lambda next_launch: next_launch.name,
+        attributes_fn=lambda next_launch: {
+            ATTR_LAUNCH_PROVIDER: next_launch.launch_service_provider.name,
+            ATTR_LAUNCH_PAD: next_launch.pad.name,
+            ATTR_LAUNCH_FACILITY: next_launch.pad.location.name,
+            ATTR_LAUNCH_PAD_COUNTRY_CODE: next_launch.pad.location.country_code,
+        },
+    ),
+    NextLaunchSensorEntityDescription(
+        key=LAUNCH_TIME,
+        icon="mdi:clock-outline",
+        name="Launch time",
+        device_class=SensorDeviceClass.TIMESTAMP,
+        value_fn=lambda next_launch: parse_datetime(next_launch.net),
+        attributes_fn=lambda next_launch: {
+            ATTR_WINDOW_START: next_launch.window_start,
+            ATTR_WINDOW_END: next_launch.window_end,
+            ATTR_STREAM_LIVE: next_launch.webcast_live,
+        },
+    ),
+    NextLaunchSensorEntityDescription(
+        key=LAUNCH_PROBABILITY,
+        icon="mdi:dice-multiple",
+        name="Launch Probability",
+        native_unit_of_measurement=PERCENTAGE,
+        value_fn=lambda next_launch: next_launch.probability
+        if next_launch.probability != -1
+        else STATE_UNKNOWN,
+        attributes_fn=lambda next_launch: None,
+    ),
 )
 
 
@@ -76,65 +133,59 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the sensor platform."""
-
     name = entry.data.get(CONF_NAME, DEFAULT_NAME)
-
     coordinator = hass.data[DOMAIN]
 
     async_add_entities(
-        [
-            NextLaunchSensor(
-                coordinator,
-                entry.entry_id,
-                name,
-                SENSOR_DESCRIPTION,
-            ),
-        ]
+        NextLaunchSensor(
+            coordinator=coordinator,
+            entry_id=entry.entry_id,
+            description=description,
+            name=name if description.key == NEXT_LAUNCH else None,
+        )
+        for description in SENSOR_DESCRIPTIONS
     )
 
 
 class NextLaunchSensor(CoordinatorEntity, SensorEntity):
-    """Representation of the next launch sensor."""
+    """Representation of the next launch sensors."""
 
     _attr_attribution = ATTRIBUTION
     _next_launch: Launch | None = None
+    entity_description: NextLaunchSensorEntityDescription
 
     def __init__(
         self,
         coordinator: DataUpdateCoordinator,
         entry_id: str,
-        name: str,
-        description: SensorEntityDescription,
+        description: NextLaunchSensorEntityDescription,
+        name: str | None = None,
     ) -> None:
-        """Initialize a Launch Library entity."""
+        """Initialize a Launch Library sensor."""
         super().__init__(coordinator)
-        self._attr_name = name
+        if name:
+            self._attr_name = name
         self._attr_unique_id = f"{entry_id}_{description.key}"
         self.entity_description = description
 
     @property
-    def available(self) -> bool:
-        """Return if the sensor is available."""
-        return super().available and self._next_launch is not None
-
-    @property
-    def native_value(self) -> str | None:
+    def native_value(self) -> datetime | str | int | None:
         """Return the state of the sensor."""
         if self._next_launch is None:
             return None
-        return self._next_launch.name
+        return self.entity_description.value_fn(self._next_launch)
 
     @property
     def extra_state_attributes(self) -> dict[str, Any] | None:
         """Return the attributes of the sensor."""
         if self._next_launch is None:
             return None
-        return {
-            ATTR_LAUNCH_TIME: self._next_launch.net,
-            ATTR_AGENCY: self._next_launch.launch_service_provider.name,
-            ATTR_AGENCY_COUNTRY_CODE: self._next_launch.pad.location.country_code,
-            ATTR_STREAM: self._next_launch.webcast_live,
-        }
+        return self.entity_description.attributes_fn(self._next_launch)
+
+    @property
+    def available(self) -> bool:
+        """Return if the sensor is available."""
+        return super().available and self._next_launch is not None
 
     @callback
     def _handle_coordinator_update(self) -> None: