Safe enum conversion (#37)

This commit is contained in:
Otto Winter 2021-06-18 16:57:07 +02:00 committed by GitHub
parent ee2c5d9152
commit a72957e2f0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 47 additions and 37 deletions

View File

@ -1,5 +1,5 @@
import enum import enum
from typing import List, Dict from typing import List, Dict, TypeVar, Optional, Type
import attr import attr
@ -10,6 +10,28 @@ import attr
# for a field (False, 0, empty string, enum with value 0, ...) # 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 @attr.s
class APIVersion: class APIVersion:
major = attr.ib(type=int, default=0) major = attr.ib(type=int, default=0)
@ -62,18 +84,18 @@ class CoverInfo(EntityInfo):
device_class = attr.ib(type=str, default='') device_class = attr.ib(type=str, default='')
class LegacyCoverState(enum.IntEnum): class LegacyCoverState(APIIntEnum):
OPEN = 0 OPEN = 0
CLOSED = 1 CLOSED = 1
class LegacyCoverCommand(enum.IntEnum): class LegacyCoverCommand(APIIntEnum):
OPEN = 0 OPEN = 0
CLOSE = 1 CLOSE = 1
STOP = 2 STOP = 2
class CoverOperation(enum.IntEnum): class CoverOperation(APIIntEnum):
IDLE = 0 IDLE = 0
IS_OPENING = 1 IS_OPENING = 1
IS_CLOSING = 2 IS_CLOSING = 2
@ -81,11 +103,11 @@ class CoverOperation(enum.IntEnum):
@attr.s @attr.s
class CoverState(EntityState): class CoverState(EntityState):
legacy_state = attr.ib(type=LegacyCoverState, converter=LegacyCoverState, legacy_state = attr.ib(type=Optional[LegacyCoverState], converter=LegacyCoverState.convert,
default=LegacyCoverState.OPEN) default=LegacyCoverState.OPEN)
position = attr.ib(type=float, default=0.0) position = attr.ib(type=float, default=0.0)
tilt = 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) default=CoverOperation.IDLE)
def is_closed(self, api_version: APIVersion): def is_closed(self, api_version: APIVersion):
@ -103,13 +125,13 @@ class FanInfo(EntityInfo):
supported_speed_levels = attr.ib(type=int, default=0) supported_speed_levels = attr.ib(type=int, default=0)
class FanSpeed(enum.IntEnum): class FanSpeed(APIIntEnum):
LOW = 0 LOW = 0
MEDIUM = 1 MEDIUM = 1
HIGH = 2 HIGH = 2
class FanDirection(enum.IntEnum): class FanDirection(APIIntEnum):
FORWARD = 0 FORWARD = 0
REVERSE = 1 REVERSE = 1
@ -118,9 +140,9 @@ class FanDirection(enum.IntEnum):
class FanState(EntityState): class FanState(EntityState):
state = attr.ib(type=bool, default=False) state = attr.ib(type=bool, default=False)
oscillating = 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) 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 ==================== # ==================== LIGHT ====================
@ -148,7 +170,7 @@ class LightState(EntityState):
# ==================== SENSOR ==================== # ==================== SENSOR ====================
class SensorStateClass(enum.IntEnum): class SensorStateClass(APIIntEnum):
NONE = 0 NONE = 0
MEASUREMENT = 1 MEASUREMENT = 1
@ -159,7 +181,7 @@ class SensorInfo(EntityInfo):
unit_of_measurement = attr.ib(type=str, default='') unit_of_measurement = attr.ib(type=str, default='')
accuracy_decimals = attr.ib(type=int, default=0) accuracy_decimals = attr.ib(type=int, default=0)
force_update = attr.ib(type=bool, default=False) 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 @attr.s
@ -204,7 +226,7 @@ class CameraState(EntityState):
# ==================== CLIMATE ==================== # ==================== CLIMATE ====================
class ClimateMode(enum.IntEnum): class ClimateMode(APIIntEnum):
OFF = 0 OFF = 0
AUTO = 1 AUTO = 1
COOL = 2 COOL = 2
@ -213,7 +235,7 @@ class ClimateMode(enum.IntEnum):
DRY = 5 DRY = 5
class ClimateFanMode(enum.IntEnum): class ClimateFanMode(APIIntEnum):
ON = 0 ON = 0
OFF = 1 OFF = 1
AUTO = 2 AUTO = 2
@ -225,14 +247,14 @@ class ClimateFanMode(enum.IntEnum):
DIFFUSE = 8 DIFFUSE = 8
class ClimateSwingMode(enum.IntEnum): class ClimateSwingMode(APIIntEnum):
OFF = 0 OFF = 0
BOTH = 1 BOTH = 1
VERTICAL = 2 VERTICAL = 2
HORIZONTAL = 3 HORIZONTAL = 3
class ClimateAction(enum.IntEnum): class ClimateAction(APIIntEnum):
OFF = 0 OFF = 0
COOLING = 2 COOLING = 2
HEATING = 3 HEATING = 3
@ -241,23 +263,11 @@ class ClimateAction(enum.IntEnum):
FAN = 6 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 @attr.s
class ClimateInfo(EntityInfo): class ClimateInfo(EntityInfo):
supports_current_temperature = attr.ib(type=bool, default=False) supports_current_temperature = attr.ib(type=bool, default=False)
supports_two_point_target_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) factory=list)
visual_min_temperature = attr.ib(type=float, default=0.0) visual_min_temperature = attr.ib(type=float, default=0.0)
visual_max_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_away = attr.ib(type=bool, default=False)
supports_action = attr.ib(type=bool, default=False) supports_action = attr.ib(type=bool, default=False)
supported_fan_modes = attr.ib( 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( 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 @attr.s
class ClimateState(EntityState): class ClimateState(EntityState):
mode = attr.ib(type=ClimateMode, converter=ClimateMode, mode = attr.ib(type=Optional[ClimateMode], converter=ClimateMode.convert,
default=ClimateMode.OFF) default=ClimateMode.OFF)
action = attr.ib(type=ClimateAction, converter=ClimateAction, action = attr.ib(type=Optional[ClimateAction], converter=ClimateAction.convert,
default=ClimateAction.OFF) default=ClimateAction.OFF)
current_temperature = attr.ib(type=float, default=0.0) current_temperature = attr.ib(type=float, default=0.0)
target_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) target_temperature_high = attr.ib(type=float, default=0.0)
away = attr.ib(type=bool, default=False) away = attr.ib(type=bool, default=False)
fan_mode = attr.ib( fan_mode = attr.ib(
type=ClimateFanMode, converter=ClimateFanMode, default=ClimateFanMode.ON type=Optional[ClimateFanMode], converter=ClimateFanMode.convert, default=ClimateFanMode.ON
) )
swing_mode = attr.ib( 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) factory=dict)
class UserServiceArgType(enum.IntEnum): class UserServiceArgType(APIIntEnum):
BOOL = 0 BOOL = 0
INT = 1 INT = 1
FLOAT = 2 FLOAT = 2
@ -339,7 +349,7 @@ def _attr_obj_from_dict(cls, **kwargs):
@attr.s @attr.s
class UserServiceArg: class UserServiceArg:
name = attr.ib(type=str, default='') 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) default=UserServiceArgType.BOOL)