ESPHome rework EsphomeEnumMapper for safe enum mappings (#51975)

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):

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)
         )

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:

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"],

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):

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

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