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