From 03ec7b3d0b4481912fb491c7c3c79080de6d0116 Mon Sep 17 00:00:00 2001 From: Otto Winter Date: Tue, 22 Jun 2021 06:22:38 +0200 Subject: [PATCH] ESPHome rework EsphomeEnumMapper for safe enum mappings (#51975) --- homeassistant/components/esphome/__init__.py | 30 ++++------ homeassistant/components/esphome/climate.py | 58 +++++++++---------- homeassistant/components/esphome/fan.py | 16 ++--- .../components/esphome/manifest.json | 2 +- homeassistant/components/esphome/sensor.py | 10 ++-- requirements_all.txt | 2 +- requirements_test_all.txt | 2 +- 7 files changed, 54 insertions(+), 66 deletions(-) diff --git a/homeassistant/components/esphome/__init__.py b/homeassistant/components/esphome/__init__.py index 607af8cc47d..7783047a662 100644 --- a/homeassistant/components/esphome/__init__.py +++ b/homeassistant/components/esphome/__init__.py @@ -5,7 +5,7 @@ import asyncio import functools import logging import math -from typing import Callable +from typing import Generic, TypeVar from aioesphomeapi import ( APIClient, @@ -48,6 +48,7 @@ from .entry_data import RuntimeEntryData DOMAIN = "esphome" _LOGGER = logging.getLogger(__name__) +_T = TypeVar("_T") STORAGE_VERSION = 1 @@ -721,30 +722,23 @@ def esphome_state_property(func): return _wrapper -class EsphomeEnumMapper: +class EsphomeEnumMapper(Generic[_T]): """Helper class to convert between hass and esphome enum values.""" - def __init__(self, func: Callable[[], dict[int, str]]) -> None: + def __init__(self, mapping: dict[_T, str]) -> None: """Construct a EsphomeEnumMapper.""" - self._func = func + # Add none mapping + mapping = {None: None, **mapping} + self._mapping = mapping + self._inverse: dict[str, _T] = {v: k for k, v in mapping.items()} - def from_esphome(self, value: int) -> str: + def from_esphome(self, value: _T | None) -> str | None: """Convert from an esphome int representation to a hass string.""" - return self._func()[value] + return self._mapping[value] - def from_hass(self, value: str) -> int: + def from_hass(self, value: str) -> _T: """Convert from a hass string to a esphome int representation.""" - inverse = {v: k for k, v in self._func().items()} - return inverse[value] - - -def esphome_map_enum(func: Callable[[], dict[int, str]]): - """Map esphome int enum values to hass string constants. - - This class has to be used as a decorator. This ensures the aioesphomeapi - import is only happening at runtime. - """ - return EsphomeEnumMapper(func) + return self._inverse[value] class EsphomeBaseEntity(Entity): diff --git a/homeassistant/components/esphome/climate.py b/homeassistant/components/esphome/climate.py index 5d21d495ec2..15edcdd8150 100644 --- a/homeassistant/components/esphome/climate.py +++ b/homeassistant/components/esphome/climate.py @@ -58,7 +58,7 @@ from homeassistant.const import ( from . import ( EsphomeEntity, - esphome_map_enum, + EsphomeEnumMapper, esphome_state_property, platform_async_setup_entry, ) @@ -77,9 +77,8 @@ async def async_setup_entry(hass, entry, async_add_entities): ) -@esphome_map_enum -def _climate_modes(): - return { +_CLIMATE_MODES: EsphomeEnumMapper[ClimateMode] = EsphomeEnumMapper( + { ClimateMode.OFF: HVAC_MODE_OFF, ClimateMode.AUTO: HVAC_MODE_HEAT_COOL, ClimateMode.COOL: HVAC_MODE_COOL, @@ -87,11 +86,9 @@ def _climate_modes(): ClimateMode.FAN_ONLY: HVAC_MODE_FAN_ONLY, ClimateMode.DRY: HVAC_MODE_DRY, } - - -@esphome_map_enum -def _climate_actions(): - return { +) +_CLIMATE_ACTIONS: EsphomeEnumMapper[ClimateAction] = EsphomeEnumMapper( + { ClimateAction.OFF: CURRENT_HVAC_OFF, ClimateAction.COOLING: CURRENT_HVAC_COOL, ClimateAction.HEATING: CURRENT_HVAC_HEAT, @@ -99,11 +96,9 @@ def _climate_actions(): ClimateAction.DRYING: CURRENT_HVAC_DRY, ClimateAction.FAN: CURRENT_HVAC_FAN, } - - -@esphome_map_enum -def _fan_modes(): - return { +) +_FAN_MODES: EsphomeEnumMapper[ClimateFanMode] = EsphomeEnumMapper( + { ClimateFanMode.ON: FAN_ON, ClimateFanMode.OFF: FAN_OFF, ClimateFanMode.AUTO: FAN_AUTO, @@ -114,16 +109,15 @@ def _fan_modes(): ClimateFanMode.FOCUS: FAN_FOCUS, ClimateFanMode.DIFFUSE: FAN_DIFFUSE, } - - -@esphome_map_enum -def _swing_modes(): - return { +) +_SWING_MODES: EsphomeEnumMapper[ClimateSwingMode] = EsphomeEnumMapper( + { ClimateSwingMode.OFF: SWING_OFF, ClimateSwingMode.BOTH: SWING_BOTH, ClimateSwingMode.VERTICAL: SWING_VERTICAL, ClimateSwingMode.HORIZONTAL: SWING_HORIZONTAL, } +) class EsphomeClimateEntity(EsphomeEntity, ClimateEntity): @@ -156,7 +150,7 @@ class EsphomeClimateEntity(EsphomeEntity, ClimateEntity): def hvac_modes(self) -> list[str]: """Return the list of available operation modes.""" return [ - _climate_modes.from_esphome(mode) + _CLIMATE_MODES.from_esphome(mode) for mode in self._static_info.supported_modes ] @@ -164,7 +158,7 @@ class EsphomeClimateEntity(EsphomeEntity, ClimateEntity): def fan_modes(self): """Return the list of available fan modes.""" return [ - _fan_modes.from_esphome(mode) + _FAN_MODES.from_esphome(mode) for mode in self._static_info.supported_fan_modes ] @@ -177,7 +171,7 @@ class EsphomeClimateEntity(EsphomeEntity, ClimateEntity): def swing_modes(self): """Return the list of available swing modes.""" return [ - _swing_modes.from_esphome(mode) + _SWING_MODES.from_esphome(mode) for mode in self._static_info.supported_swing_modes ] @@ -219,7 +213,7 @@ class EsphomeClimateEntity(EsphomeEntity, ClimateEntity): @esphome_state_property def hvac_mode(self) -> str | None: """Return current operation ie. heat, cool, idle.""" - return _climate_modes.from_esphome(self._state.mode) + return _CLIMATE_MODES.from_esphome(self._state.mode) @esphome_state_property def hvac_action(self) -> str | None: @@ -227,12 +221,12 @@ class EsphomeClimateEntity(EsphomeEntity, ClimateEntity): # HA has no support feature field for hvac_action if not self._static_info.supports_action: return None - return _climate_actions.from_esphome(self._state.action) + return _CLIMATE_ACTIONS.from_esphome(self._state.action) @esphome_state_property - def fan_mode(self): + def fan_mode(self) -> str | None: """Return current fan setting.""" - return _fan_modes.from_esphome(self._state.fan_mode) + return _FAN_MODES.from_esphome(self._state.fan_mode) @esphome_state_property def preset_mode(self): @@ -240,9 +234,9 @@ class EsphomeClimateEntity(EsphomeEntity, ClimateEntity): return PRESET_AWAY if self._state.away else PRESET_HOME @esphome_state_property - def swing_mode(self): + def swing_mode(self) -> str | None: """Return current swing mode.""" - return _swing_modes.from_esphome(self._state.swing_mode) + return _SWING_MODES.from_esphome(self._state.swing_mode) @esphome_state_property def current_temperature(self) -> float | None: @@ -268,7 +262,7 @@ class EsphomeClimateEntity(EsphomeEntity, ClimateEntity): """Set new target temperature (and operation mode if set).""" data = {"key": self._static_info.key} if ATTR_HVAC_MODE in kwargs: - data["mode"] = _climate_modes.from_hass(kwargs[ATTR_HVAC_MODE]) + data["mode"] = _CLIMATE_MODES.from_hass(kwargs[ATTR_HVAC_MODE]) if ATTR_TEMPERATURE in kwargs: data["target_temperature"] = kwargs[ATTR_TEMPERATURE] if ATTR_TARGET_TEMP_LOW in kwargs: @@ -280,7 +274,7 @@ class EsphomeClimateEntity(EsphomeEntity, ClimateEntity): async def async_set_hvac_mode(self, hvac_mode: str) -> None: """Set new target operation mode.""" await self._client.climate_command( - key=self._static_info.key, mode=_climate_modes.from_hass(hvac_mode) + key=self._static_info.key, mode=_CLIMATE_MODES.from_hass(hvac_mode) ) async def async_set_preset_mode(self, preset_mode): @@ -291,11 +285,11 @@ class EsphomeClimateEntity(EsphomeEntity, ClimateEntity): async def async_set_fan_mode(self, fan_mode: str) -> None: """Set new fan mode.""" await self._client.climate_command( - key=self._static_info.key, fan_mode=_fan_modes.from_hass(fan_mode) + key=self._static_info.key, fan_mode=_FAN_MODES.from_hass(fan_mode) ) async def async_set_swing_mode(self, swing_mode: str) -> None: """Set new swing mode.""" await self._client.climate_command( - key=self._static_info.key, swing_mode=_swing_modes.from_hass(swing_mode) + key=self._static_info.key, swing_mode=_SWING_MODES.from_hass(swing_mode) ) diff --git a/homeassistant/components/esphome/fan.py b/homeassistant/components/esphome/fan.py index 5272cdef5f1..b73df42c8dc 100644 --- a/homeassistant/components/esphome/fan.py +++ b/homeassistant/components/esphome/fan.py @@ -24,7 +24,7 @@ from homeassistant.util.percentage import ( from . import ( EsphomeEntity, - esphome_map_enum, + EsphomeEnumMapper, esphome_state_property, platform_async_setup_entry, ) @@ -47,12 +47,12 @@ async def async_setup_entry( ) -@esphome_map_enum -def _fan_directions(): - return { +_FAN_DIRECTIONS: EsphomeEnumMapper[FanDirection] = EsphomeEnumMapper( + { FanDirection.FORWARD: DIRECTION_FORWARD, FanDirection.REVERSE: DIRECTION_REVERSE, } +) class EsphomeFan(EsphomeEntity, FanEntity): @@ -115,7 +115,7 @@ class EsphomeFan(EsphomeEntity, FanEntity): async def async_set_direction(self, direction: str): """Set direction of the fan.""" await self._client.fan_command( - key=self._static_info.key, direction=_fan_directions.from_hass(direction) + key=self._static_info.key, direction=_FAN_DIRECTIONS.from_hass(direction) ) # https://github.com/PyCQA/pylint/issues/3150 for all @esphome_state_property @@ -149,18 +149,18 @@ class EsphomeFan(EsphomeEntity, FanEntity): return self._static_info.supported_speed_levels @esphome_state_property - def oscillating(self) -> None: + def oscillating(self) -> bool | None: """Return the oscillation state.""" if not self._static_info.supports_oscillation: return None return self._state.oscillating @esphome_state_property - def current_direction(self) -> None: + def current_direction(self) -> str | None: """Return the current fan direction.""" if not self._static_info.supports_direction: return None - return _fan_directions.from_esphome(self._state.direction) + return _FAN_DIRECTIONS.from_esphome(self._state.direction) @property def supported_features(self) -> int: diff --git a/homeassistant/components/esphome/manifest.json b/homeassistant/components/esphome/manifest.json index 592ca616d04..b732713335e 100644 --- a/homeassistant/components/esphome/manifest.json +++ b/homeassistant/components/esphome/manifest.json @@ -3,7 +3,7 @@ "name": "ESPHome", "config_flow": true, "documentation": "https://www.home-assistant.io/integrations/esphome", - "requirements": ["aioesphomeapi==2.8.0"], + "requirements": ["aioesphomeapi==2.9.0"], "zeroconf": ["_esphomelib._tcp.local."], "codeowners": ["@OttoWinter"], "after_dependencies": ["zeroconf", "tag"], diff --git a/homeassistant/components/esphome/sensor.py b/homeassistant/components/esphome/sensor.py index 7b905aad148..d3dce2dea1b 100644 --- a/homeassistant/components/esphome/sensor.py +++ b/homeassistant/components/esphome/sensor.py @@ -25,7 +25,7 @@ from homeassistant.util import dt from . import ( EsphomeEntity, - esphome_map_enum, + EsphomeEnumMapper, esphome_state_property, platform_async_setup_entry, ) @@ -61,12 +61,12 @@ async def async_setup_entry( # pylint: disable=invalid-overridden-method -@esphome_map_enum -def _state_classes(): - return { +_STATE_CLASSES: EsphomeEnumMapper[SensorStateClass] = EsphomeEnumMapper( + { SensorStateClass.NONE: None, SensorStateClass.MEASUREMENT: STATE_CLASS_MEASUREMENT, } +) class EsphomeSensor(EsphomeEntity, SensorEntity): @@ -122,7 +122,7 @@ class EsphomeSensor(EsphomeEntity, SensorEntity): """Return the state class of this entity.""" if not self._static_info.state_class: return None - return _state_classes.from_esphome(self._static_info.state_class) + return _STATE_CLASSES.from_esphome(self._static_info.state_class) class EsphomeTextSensor(EsphomeEntity, SensorEntity): diff --git a/requirements_all.txt b/requirements_all.txt index 9c6c1e84967..5d19bd04e9d 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -160,7 +160,7 @@ aioeafm==0.1.2 aioemonitor==1.0.5 # homeassistant.components.esphome -aioesphomeapi==2.8.0 +aioesphomeapi==2.9.0 # homeassistant.components.flo aioflo==0.4.1 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index c7d6c608255..b8ff4863197 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -100,7 +100,7 @@ aioeafm==0.1.2 aioemonitor==1.0.5 # homeassistant.components.esphome -aioesphomeapi==2.8.0 +aioesphomeapi==2.9.0 # homeassistant.components.flo aioflo==0.4.1