Skip to content

Commit 7c6a39e

Browse files
epenetepenet
authored
Add wrapper class for enum values in Tuya models (home-assistant#155847)
1 parent 57c3a5c commit 7c6a39e

File tree

3 files changed

+197
-26
lines changed

3 files changed

+197
-26
lines changed

‎homeassistant/components/tuya/models.py‎

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,14 +6,70 @@
66
from dataclasses import dataclass
77
import json
88
import struct
9-
from typing import Literal, Self, overload
9+
from typing import Any, Literal, Self, overload
1010

1111
from tuya_sharing import CustomerDevice
1212

1313
from .const import DPCode, DPType
1414
from .util import remap_value
1515

1616

17+
@dataclass
18+
class DPCodeWrapper:
19+
"""Base DPCode wrapper.
20+
21+
Used as a common interface for referring to a DPCode, and
22+
access read conversion routines.
23+
"""
24+
25+
dpcode: str
26+
27+
def _read_device_status_raw(self, device: CustomerDevice) -> Any | None:
28+
"""Read the raw device status for the DPCode.
29+
30+
Private helper method for `read_device_status`.
31+
"""
32+
return device.status.get(self.dpcode)
33+
34+
def read_device_status(self, device: CustomerDevice) -> Any | None:
35+
"""Read the device value for the dpcode."""
36+
raise NotImplementedError("read_device_value must be implemented")
37+
38+
39+
@dataclass(kw_only=True)
40+
class DPCodeEnumWrapper(DPCodeWrapper):
41+
"""Simple wrapper for EnumTypeData values."""
42+
43+
enum_type_information: EnumTypeData
44+
45+
def read_device_status(self, device: CustomerDevice) -> str | None:
46+
"""Read the device value for the dpcode.
47+
48+
Values outside of the list defined by the Enum type information will
49+
return None.
50+
"""
51+
if (
52+
raw_value := self._read_device_status_raw(device)
53+
) in self.enum_type_information.range:
54+
return raw_value
55+
return None
56+
57+
@classmethod
58+
def find_dpcode(
59+
cls,
60+
device: CustomerDevice,
61+
dpcodes: str | DPCode | tuple[DPCode, ...],
62+
*,
63+
prefer_function: bool = False,
64+
) -> Self | None:
65+
"""Find and return a DPCodeEnumWrapper for the given DP codes."""
66+
if enum_type := find_dpcode(
67+
device, dpcodes, dptype=DPType.ENUM, prefer_function=prefer_function
68+
):
69+
return cls(dpcode=enum_type.dpcode, enum_type_information=enum_type)
70+
return None
71+
72+
1773
@overload
1874
def find_dpcode(
1975
device: CustomerDevice,

‎homeassistant/components/tuya/select.py‎

Lines changed: 15 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,9 @@
1111
from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
1212

1313
from . import TuyaConfigEntry
14-
from .const import TUYA_DISCOVERY_NEW, DeviceCategory, DPCode, DPType
14+
from .const import TUYA_DISCOVERY_NEW, DeviceCategory, DPCode
1515
from .entity import TuyaEntity
16-
from .models import find_dpcode
16+
from .models import DPCodeEnumWrapper
1717

1818
# All descriptions can be found here. Mostly the Enum data types in the
1919
# default instructions set of each category end up being a select.
@@ -360,9 +360,15 @@ def async_discover_device(device_ids: list[str]) -> None:
360360
device = manager.device_map[device_id]
361361
if descriptions := SELECTS.get(device.category):
362362
entities.extend(
363-
TuyaSelectEntity(device, manager, description)
363+
TuyaSelectEntity(
364+
device, manager, description, dpcode_wrapper=dpcode_wrapper
365+
)
364366
for description in descriptions
365-
if description.key in device.status
367+
if (
368+
dpcode_wrapper := DPCodeEnumWrapper.find_dpcode(
369+
device, description.key, prefer_function=True
370+
)
371+
)
366372
)
367373

368374
async_add_entities(entities)
@@ -382,35 +388,20 @@ def __init__(
382388
device: CustomerDevice,
383389
device_manager: Manager,
384390
description: SelectEntityDescription,
391+
dpcode_wrapper: DPCodeEnumWrapper,
385392
) -> None:
386393
"""Init Tuya sensor."""
387394
super().__init__(device, device_manager)
388395
self.entity_description = description
389396
self._attr_unique_id = f"{super().unique_id}{description.key}"
390-
391-
self._attr_options: list[str] = []
392-
if enum_type := find_dpcode(
393-
self.device, description.key, dptype=DPType.ENUM, prefer_function=True
394-
):
395-
self._attr_options = enum_type.range
397+
self._dpcode_wrapper = dpcode_wrapper
398+
self._attr_options = dpcode_wrapper.enum_type_information.range
396399

397400
@property
398401
def current_option(self) -> str | None:
399402
"""Return the selected entity option to represent the entity state."""
400-
# Raw value
401-
value = self.device.status.get(self.entity_description.key)
402-
if value is None or value not in self._attr_options:
403-
return None
404-
405-
return value
403+
return self._dpcode_wrapper.read_device_status(self.device)
406404

407405
def select_option(self, option: str) -> None:
408406
"""Change the selected option."""
409-
self._send_command(
410-
[
411-
{
412-
"code": self.entity_description.key,
413-
"value": option,
414-
}
415-
]
416-
)
407+
self._send_command([{"code": self._dpcode_wrapper.dpcode, "value": option}])

‎tests/components/tuya/snapshots/test_select.ambr‎

Lines changed: 125 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2928,7 +2928,7 @@
29282928
'suggested_object_id': None,
29292929
'supported_features': 0,
29302930
'translation_key': 'countdown',
2931-
'unique_id': 'tuya.owozxdzgbibizu4sjkcountdown_set',
2931+
'unique_id': 'tuya.owozxdzgbibizu4sjkcountdown',
29322932
'unit_of_measurement': None,
29332933
})
29342934
# ---
@@ -2953,6 +2953,71 @@
29532953
'state': 'unknown',
29542954
})
29552955
# ---
2956+
# name: test_platform_setup_and_discovery[select.ion1000pro_countdown_2-entry]
2957+
EntityRegistryEntrySnapshot({
2958+
'aliases': set({
2959+
}),
2960+
'area_id': None,
2961+
'capabilities': dict({
2962+
'options': list([
2963+
'1',
2964+
'2',
2965+
'3',
2966+
'4',
2967+
'5',
2968+
'6',
2969+
]),
2970+
}),
2971+
'config_entry_id': <ANY>,
2972+
'config_subentry_id': <ANY>,
2973+
'device_class': None,
2974+
'device_id': <ANY>,
2975+
'disabled_by': None,
2976+
'domain': 'select',
2977+
'entity_category': <EntityCategory.CONFIG: 'config'>,
2978+
'entity_id': 'select.ion1000pro_countdown_2',
2979+
'has_entity_name': True,
2980+
'hidden_by': None,
2981+
'icon': None,
2982+
'id': <ANY>,
2983+
'labels': set({
2984+
}),
2985+
'name': None,
2986+
'options': dict({
2987+
}),
2988+
'original_device_class': None,
2989+
'original_icon': None,
2990+
'original_name': 'Countdown',
2991+
'platform': 'tuya',
2992+
'previous_unique_id': None,
2993+
'suggested_object_id': None,
2994+
'supported_features': 0,
2995+
'translation_key': 'countdown',
2996+
'unique_id': 'tuya.owozxdzgbibizu4sjkcountdown_set',
2997+
'unit_of_measurement': None,
2998+
})
2999+
# ---
3000+
# name: test_platform_setup_and_discovery[select.ion1000pro_countdown_2-state]
3001+
StateSnapshot({
3002+
'attributes': ReadOnlyDict({
3003+
'friendly_name': 'ION1000PRO Countdown',
3004+
'options': list([
3005+
'1',
3006+
'2',
3007+
'3',
3008+
'4',
3009+
'5',
3010+
'6',
3011+
]),
3012+
}),
3013+
'context': <ANY>,
3014+
'entity_id': 'select.ion1000pro_countdown_2',
3015+
'last_changed': <ANY>,
3016+
'last_reported': <ANY>,
3017+
'last_updated': <ANY>,
3018+
'state': 'unknown',
3019+
})
3020+
# ---
29563021
# name: test_platform_setup_and_discovery[select.jardin_fraises_power_on_behavior-entry]
29573022
EntityRegistryEntrySnapshot({
29583023
'aliases': set({
@@ -4088,6 +4153,65 @@
40884153
'state': 'power_on',
40894154
})
40904155
# ---
4156+
# name: test_platform_setup_and_discovery[select.seating_side_6_ch_smart_switch_power_on_behavior-entry]
4157+
EntityRegistryEntrySnapshot({
4158+
'aliases': set({
4159+
}),
4160+
'area_id': None,
4161+
'capabilities': dict({
4162+
'options': list([
4163+
'0',
4164+
'1',
4165+
'2',
4166+
]),
4167+
}),
4168+
'config_entry_id': <ANY>,
4169+
'config_subentry_id': <ANY>,
4170+
'device_class': None,
4171+
'device_id': <ANY>,
4172+
'disabled_by': None,
4173+
'domain': 'select',
4174+
'entity_category': <EntityCategory.CONFIG: 'config'>,
4175+
'entity_id': 'select.seating_side_6_ch_smart_switch_power_on_behavior',
4176+
'has_entity_name': True,
4177+
'hidden_by': None,
4178+
'icon': None,
4179+
'id': <ANY>,
4180+
'labels': set({
4181+
}),
4182+
'name': None,
4183+
'options': dict({
4184+
}),
4185+
'original_device_class': None,
4186+
'original_icon': None,
4187+
'original_name': 'Power on behavior',
4188+
'platform': 'tuya',
4189+
'previous_unique_id': None,
4190+
'suggested_object_id': None,
4191+
'supported_features': 0,
4192+
'translation_key': 'relay_status',
4193+
'unique_id': 'tuya.kxxrbv93k2vvkconqdtrelay_status',
4194+
'unit_of_measurement': None,
4195+
})
4196+
# ---
4197+
# name: test_platform_setup_and_discovery[select.seating_side_6_ch_smart_switch_power_on_behavior-state]
4198+
StateSnapshot({
4199+
'attributes': ReadOnlyDict({
4200+
'friendly_name': 'Seating side 6-ch Smart Switch Power on behavior',
4201+
'options': list([
4202+
'0',
4203+
'1',
4204+
'2',
4205+
]),
4206+
}),
4207+
'context': <ANY>,
4208+
'entity_id': 'select.seating_side_6_ch_smart_switch_power_on_behavior',
4209+
'last_changed': <ANY>,
4210+
'last_reported': <ANY>,
4211+
'last_updated': <ANY>,
4212+
'state': 'unknown',
4213+
})
4214+
# ---
40914215
# name: test_platform_setup_and_discovery[select.security_light_indicator_light_mode-entry]
40924216
EntityRegistryEntrySnapshot({
40934217
'aliases': set({

0 commit comments

Comments
 (0)