Add generic typing for powerwall sensors (#109008)

pull/109186/head
Josh Pettersen 2024-01-30 18:28:27 -08:00 committed by GitHub
parent 320bf53f75
commit 2f9f0bae46
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 17 additions and 11 deletions

View File

@ -3,7 +3,7 @@ from __future__ import annotations
from collections.abc import Callable from collections.abc import Callable
from dataclasses import dataclass from dataclasses import dataclass
from typing import TYPE_CHECKING from typing import TYPE_CHECKING, Generic, TypeVar
from tesla_powerwall import MeterResponse, MeterType from tesla_powerwall import MeterResponse, MeterType
@ -23,6 +23,7 @@ from homeassistant.const import (
UnitOfPower, UnitOfPower,
) )
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity import Entity
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
from .const import DOMAIN, POWERWALL_COORDINATOR from .const import DOMAIN, POWERWALL_COORDINATOR
@ -32,17 +33,22 @@ from .models import PowerwallRuntimeData
_METER_DIRECTION_EXPORT = "export" _METER_DIRECTION_EXPORT = "export"
_METER_DIRECTION_IMPORT = "import" _METER_DIRECTION_IMPORT = "import"
_ValueParamT = TypeVar("_ValueParamT")
_ValueT = TypeVar("_ValueT", bound=float)
@dataclass(frozen=True) @dataclass(frozen=True)
class PowerwallRequiredKeysMixin: class PowerwallRequiredKeysMixin(Generic[_ValueParamT, _ValueT]):
"""Mixin for required keys.""" """Mixin for required keys."""
value_fn: Callable[[MeterResponse], float] value_fn: Callable[[_ValueParamT], _ValueT]
@dataclass(frozen=True) @dataclass(frozen=True)
class PowerwallSensorEntityDescription( class PowerwallSensorEntityDescription(
SensorEntityDescription, PowerwallRequiredKeysMixin SensorEntityDescription,
PowerwallRequiredKeysMixin[_ValueParamT, _ValueT],
Generic[_ValueParamT, _ValueT],
): ):
"""Describes Powerwall entity.""" """Describes Powerwall entity."""
@ -68,7 +74,7 @@ def _get_meter_average_voltage(meter: MeterResponse) -> float:
POWERWALL_INSTANT_SENSORS = ( POWERWALL_INSTANT_SENSORS = (
PowerwallSensorEntityDescription( PowerwallSensorEntityDescription[MeterResponse, float](
key="instant_power", key="instant_power",
translation_key="instant_power", translation_key="instant_power",
state_class=SensorStateClass.MEASUREMENT, state_class=SensorStateClass.MEASUREMENT,
@ -76,7 +82,7 @@ POWERWALL_INSTANT_SENSORS = (
native_unit_of_measurement=UnitOfPower.KILO_WATT, native_unit_of_measurement=UnitOfPower.KILO_WATT,
value_fn=_get_meter_power, value_fn=_get_meter_power,
), ),
PowerwallSensorEntityDescription( PowerwallSensorEntityDescription[MeterResponse, float](
key="instant_frequency", key="instant_frequency",
translation_key="instant_frequency", translation_key="instant_frequency",
state_class=SensorStateClass.MEASUREMENT, state_class=SensorStateClass.MEASUREMENT,
@ -85,7 +91,7 @@ POWERWALL_INSTANT_SENSORS = (
entity_registry_enabled_default=False, entity_registry_enabled_default=False,
value_fn=_get_meter_frequency, value_fn=_get_meter_frequency,
), ),
PowerwallSensorEntityDescription( PowerwallSensorEntityDescription[MeterResponse, float](
key="instant_current", key="instant_current",
translation_key="instant_current", translation_key="instant_current",
state_class=SensorStateClass.MEASUREMENT, state_class=SensorStateClass.MEASUREMENT,
@ -94,7 +100,7 @@ POWERWALL_INSTANT_SENSORS = (
entity_registry_enabled_default=False, entity_registry_enabled_default=False,
value_fn=_get_meter_total_current, value_fn=_get_meter_total_current,
), ),
PowerwallSensorEntityDescription( PowerwallSensorEntityDescription[MeterResponse, float](
key="instant_voltage", key="instant_voltage",
translation_key="instant_voltage", translation_key="instant_voltage",
state_class=SensorStateClass.MEASUREMENT, state_class=SensorStateClass.MEASUREMENT,
@ -116,7 +122,7 @@ async def async_setup_entry(
coordinator = powerwall_data[POWERWALL_COORDINATOR] coordinator = powerwall_data[POWERWALL_COORDINATOR]
assert coordinator is not None assert coordinator is not None
data = coordinator.data data = coordinator.data
entities: list[PowerWallEntity] = [ entities: list[Entity] = [
PowerWallChargeSensor(powerwall_data), PowerWallChargeSensor(powerwall_data),
] ]
@ -156,13 +162,13 @@ class PowerWallChargeSensor(PowerWallEntity, SensorEntity):
class PowerWallEnergySensor(PowerWallEntity, SensorEntity): class PowerWallEnergySensor(PowerWallEntity, SensorEntity):
"""Representation of an Powerwall Energy sensor.""" """Representation of an Powerwall Energy sensor."""
entity_description: PowerwallSensorEntityDescription entity_description: PowerwallSensorEntityDescription[MeterResponse, float]
def __init__( def __init__(
self, self,
powerwall_data: PowerwallRuntimeData, powerwall_data: PowerwallRuntimeData,
meter: MeterType, meter: MeterType,
description: PowerwallSensorEntityDescription, description: PowerwallSensorEntityDescription[MeterResponse, float],
) -> None: ) -> None:
"""Initialize the sensor.""" """Initialize the sensor."""
self.entity_description = description self.entity_description = description