From a72957e2f0ec20dc5563ac8ed0407a8346c6d545 Mon Sep 17 00:00:00 2001 From: Otto Winter Date: Fri, 18 Jun 2021 16:57:07 +0200 Subject: [PATCH] Safe enum conversion (#37) --- aioesphomeapi/model.py | 84 +++++++++++++++++++++++------------------- 1 file changed, 47 insertions(+), 37 deletions(-) diff --git a/aioesphomeapi/model.py b/aioesphomeapi/model.py index 72df712..38dbac7 100644 --- a/aioesphomeapi/model.py +++ b/aioesphomeapi/model.py @@ -1,5 +1,5 @@ import enum -from typing import List, Dict +from typing import List, Dict, TypeVar, Optional, Type import attr @@ -10,6 +10,28 @@ import attr # for a field (False, 0, empty string, enum with value 0, ...) +_T = TypeVar('_T') + +class APIIntEnum(enum.IntEnum): + """Base class for int enum values in API model.""" + @classmethod + def convert(cls: Type[_T], value: int) -> Optional[_T]: + try: + return cls(value) + except ValueError: + return None + + @classmethod + def convert_list(cls: Type[_T], value: List[int]) -> List[_T]: + ret = [] + for x in value: + try: + ret.append(cls(x)) + except ValueError: + pass + return ret + + @attr.s class APIVersion: major = attr.ib(type=int, default=0) @@ -62,18 +84,18 @@ class CoverInfo(EntityInfo): device_class = attr.ib(type=str, default='') -class LegacyCoverState(enum.IntEnum): +class LegacyCoverState(APIIntEnum): OPEN = 0 CLOSED = 1 -class LegacyCoverCommand(enum.IntEnum): +class LegacyCoverCommand(APIIntEnum): OPEN = 0 CLOSE = 1 STOP = 2 -class CoverOperation(enum.IntEnum): +class CoverOperation(APIIntEnum): IDLE = 0 IS_OPENING = 1 IS_CLOSING = 2 @@ -81,11 +103,11 @@ class CoverOperation(enum.IntEnum): @attr.s class CoverState(EntityState): - legacy_state = attr.ib(type=LegacyCoverState, converter=LegacyCoverState, + legacy_state = attr.ib(type=Optional[LegacyCoverState], converter=LegacyCoverState.convert, default=LegacyCoverState.OPEN) position = attr.ib(type=float, default=0.0) tilt = attr.ib(type=float, default=0.0) - current_operation = attr.ib(type=CoverOperation, converter=CoverOperation, + current_operation = attr.ib(type=Optional[CoverOperation], converter=CoverOperation.convert, default=CoverOperation.IDLE) def is_closed(self, api_version: APIVersion): @@ -103,13 +125,13 @@ class FanInfo(EntityInfo): supported_speed_levels = attr.ib(type=int, default=0) -class FanSpeed(enum.IntEnum): +class FanSpeed(APIIntEnum): LOW = 0 MEDIUM = 1 HIGH = 2 -class FanDirection(enum.IntEnum): +class FanDirection(APIIntEnum): FORWARD = 0 REVERSE = 1 @@ -118,9 +140,9 @@ class FanDirection(enum.IntEnum): class FanState(EntityState): state = attr.ib(type=bool, default=False) oscillating = attr.ib(type=bool, default=False) - speed = attr.ib(type=FanSpeed, converter=FanSpeed, default=FanSpeed.LOW) + speed = attr.ib(type=Optional[FanSpeed], converter=FanSpeed.convert, default=FanSpeed.LOW) speed_level = attr.ib(type=int, default=0) - direction = attr.ib(type=FanDirection, converter=FanDirection, default=FanDirection.FORWARD) + direction = attr.ib(type=Optional[FanDirection], converter=FanDirection.convert, default=FanDirection.FORWARD) # ==================== LIGHT ==================== @@ -148,7 +170,7 @@ class LightState(EntityState): # ==================== SENSOR ==================== -class SensorStateClass(enum.IntEnum): +class SensorStateClass(APIIntEnum): NONE = 0 MEASUREMENT = 1 @@ -159,7 +181,7 @@ class SensorInfo(EntityInfo): unit_of_measurement = attr.ib(type=str, default='') accuracy_decimals = attr.ib(type=int, default=0) force_update = attr.ib(type=bool, default=False) - state_class = attr.ib(type=SensorStateClass, converter=SensorStateClass, default=SensorStateClass.NONE) + state_class = attr.ib(type=Optional[SensorStateClass], converter=SensorStateClass.convert, default=SensorStateClass.NONE) @attr.s @@ -204,7 +226,7 @@ class CameraState(EntityState): # ==================== CLIMATE ==================== -class ClimateMode(enum.IntEnum): +class ClimateMode(APIIntEnum): OFF = 0 AUTO = 1 COOL = 2 @@ -213,7 +235,7 @@ class ClimateMode(enum.IntEnum): DRY = 5 -class ClimateFanMode(enum.IntEnum): +class ClimateFanMode(APIIntEnum): ON = 0 OFF = 1 AUTO = 2 @@ -225,14 +247,14 @@ class ClimateFanMode(enum.IntEnum): DIFFUSE = 8 -class ClimateSwingMode(enum.IntEnum): +class ClimateSwingMode(APIIntEnum): OFF = 0 BOTH = 1 VERTICAL = 2 HORIZONTAL = 3 -class ClimateAction(enum.IntEnum): +class ClimateAction(APIIntEnum): OFF = 0 COOLING = 2 HEATING = 3 @@ -241,23 +263,11 @@ class ClimateAction(enum.IntEnum): FAN = 6 -def _convert_climate_modes(value): - return [ClimateMode(val) for val in value] - - -def _convert_climate_fan_modes(value): - return [ClimateFanMode(val) for val in value] - - -def _convert_climate_swing_modes(value): - return [ClimateSwingMode(val) for val in value] - - @attr.s class ClimateInfo(EntityInfo): supports_current_temperature = attr.ib(type=bool, default=False) supports_two_point_target_temperature = attr.ib(type=bool, default=False) - supported_modes = attr.ib(type=List[ClimateMode], converter=_convert_climate_modes, + supported_modes = attr.ib(type=List[ClimateMode], converter=ClimateMode.convert_list, factory=list) visual_min_temperature = attr.ib(type=float, default=0.0) visual_max_temperature = attr.ib(type=float, default=0.0) @@ -265,18 +275,18 @@ class ClimateInfo(EntityInfo): supports_away = attr.ib(type=bool, default=False) supports_action = attr.ib(type=bool, default=False) supported_fan_modes = attr.ib( - type=List[ClimateFanMode], converter=_convert_climate_fan_modes, factory=list + type=List[ClimateFanMode], converter=ClimateFanMode.convert_list, factory=list ) supported_swing_modes = attr.ib( - type=List[ClimateSwingMode], converter=_convert_climate_swing_modes, factory=list + type=List[ClimateSwingMode], converter=ClimateSwingMode.convert_list, factory=list ) @attr.s class ClimateState(EntityState): - mode = attr.ib(type=ClimateMode, converter=ClimateMode, + mode = attr.ib(type=Optional[ClimateMode], converter=ClimateMode.convert, default=ClimateMode.OFF) - action = attr.ib(type=ClimateAction, converter=ClimateAction, + action = attr.ib(type=Optional[ClimateAction], converter=ClimateAction.convert, default=ClimateAction.OFF) current_temperature = attr.ib(type=float, default=0.0) target_temperature = attr.ib(type=float, default=0.0) @@ -284,10 +294,10 @@ class ClimateState(EntityState): target_temperature_high = attr.ib(type=float, default=0.0) away = attr.ib(type=bool, default=False) fan_mode = attr.ib( - type=ClimateFanMode, converter=ClimateFanMode, default=ClimateFanMode.ON + type=Optional[ClimateFanMode], converter=ClimateFanMode.convert, default=ClimateFanMode.ON ) swing_mode = attr.ib( - type=ClimateSwingMode, converter=ClimateSwingMode, default=ClimateSwingMode.OFF + type=Optional[ClimateSwingMode], converter=ClimateSwingMode.convert, default=ClimateSwingMode.OFF ) @@ -321,7 +331,7 @@ class HomeassistantServiceCall: factory=dict) -class UserServiceArgType(enum.IntEnum): +class UserServiceArgType(APIIntEnum): BOOL = 0 INT = 1 FLOAT = 2 @@ -339,7 +349,7 @@ def _attr_obj_from_dict(cls, **kwargs): @attr.s class UserServiceArg: name = attr.ib(type=str, default='') - type_ = attr.ib(type=UserServiceArgType, converter=UserServiceArgType, + type_ = attr.ib(type=Optional[UserServiceArgType], converter=UserServiceArgType.convert, default=UserServiceArgType.BOOL)