From 8ec38ef034f4d09686076ae3616e7c29f06706d0 Mon Sep 17 00:00:00 2001 From: Franck Nijhof Date: Tue, 12 Oct 2021 16:49:46 +0200 Subject: [PATCH] Refactor Tuya device handling (#57545) * Refactor Tuya device handling * Tweak --- homeassistant/components/tuya/__init__.py | 104 ++++++++++------------ homeassistant/components/tuya/climate.py | 61 ++++--------- homeassistant/components/tuya/const.py | 32 +++++-- homeassistant/components/tuya/fan.py | 55 +++--------- homeassistant/components/tuya/light.py | 60 +++---------- homeassistant/components/tuya/scene.py | 42 ++++----- homeassistant/components/tuya/switch.py | 91 ++++++++----------- 7 files changed, 164 insertions(+), 281 deletions(-) diff --git a/homeassistant/components/tuya/__init__.py b/homeassistant/components/tuya/__init__.py index 8dd9979f3e7..1c7908ab986 100644 --- a/homeassistant/components/tuya/__init__.py +++ b/homeassistant/components/tuya/__init__.py @@ -1,7 +1,8 @@ """Support for Tuya Smart devices.""" +from __future__ import annotations -import itertools import logging +from typing import NamedTuple from tuya_iot import ( AuthType, @@ -30,24 +31,25 @@ from .const import ( CONF_USERNAME, DOMAIN, PLATFORMS, - TUYA_DEVICE_MANAGER, TUYA_DISCOVERY_NEW, - TUYA_HA_DEVICES, TUYA_HA_SIGNAL_UPDATE_ENTITY, - TUYA_HA_TUYA_MAP, - TUYA_HOME_MANAGER, - TUYA_MQTT_LISTENER, + TUYA_SUPPORTED_PRODUCT_CATEGORIES, ) _LOGGER = logging.getLogger(__name__) +class HomeAssistantTuyaData(NamedTuple): + """Tuya data stored in the Home Assistant data object.""" + + device_listener: TuyaDeviceListener + device_manager: TuyaDeviceManager + home_manager: TuyaHomeManager + + async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Async setup hass config entry.""" - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = { - TUYA_HA_TUYA_MAP: {}, - TUYA_HA_DEVICES: set(), - } + hass.data.setdefault(DOMAIN, {}) # Project type has been renamed to auth type in the upstream Tuya IoT SDK. # This migrates existing config entries to reflect that name change. @@ -98,25 +100,28 @@ async def _init_tuya_sdk(hass: HomeAssistant, entry: ConfigEntry) -> bool: tuya_mq = TuyaOpenMQ(api) tuya_mq.start() + device_ids: set[str] = set() device_manager = TuyaDeviceManager(api, tuya_mq) - - # Get device list home_manager = TuyaHomeManager(api, tuya_mq, device_manager) - await hass.async_add_executor_job(home_manager.update_device_cache) - hass.data[DOMAIN][entry.entry_id][TUYA_HOME_MANAGER] = home_manager - - listener = DeviceListener(hass, entry) - hass.data[DOMAIN][entry.entry_id][TUYA_MQTT_LISTENER] = listener + listener = DeviceListener(hass, device_manager, device_ids) device_manager.add_device_listener(listener) - hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER] = device_manager - # Clean up device entities + hass.data[DOMAIN][entry.entry_id] = HomeAssistantTuyaData( + device_listener=listener, + device_manager=device_manager, + home_manager=home_manager, + ) + + # Get devices & clean up device entities + await hass.async_add_executor_job(home_manager.update_device_cache) await cleanup_device_registry(hass, device_manager) - _LOGGER.debug("init support type->%s", PLATFORMS) + # Register known device IDs + for device in device_manager.device_map.values(): + if device.category in TUYA_SUPPORTED_PRODUCT_CATEGORIES: + device_ids.add(device.id) hass.config_entries.async_setup_platforms(entry, PLATFORMS) - return True @@ -134,17 +139,13 @@ async def cleanup_device_registry( async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Unloading the Tuya platforms.""" - _LOGGER.debug("integration unload") unload = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) if unload: - device_manager = hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER] - device_manager.mq.stop() - device_manager.remove_device_listener( - hass.data[DOMAIN][entry.entry_id][TUYA_MQTT_LISTENER] - ) + hass_data: HomeAssistantTuyaData = hass.data[DOMAIN][entry.entry_id] + hass_data.device_manager.mq.stop() + hass_data.device_manager.remove_device_listener(hass_data.device_listener) hass.data[DOMAIN].pop(entry.entry_id) - if not hass.data[DOMAIN]: hass.data.pop(DOMAIN) @@ -154,14 +155,20 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: class DeviceListener(TuyaDeviceListener): """Device Update Listener.""" - def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None: + def __init__( + self, + hass: HomeAssistant, + device_manager: TuyaDeviceManager, + device_ids: set[str], + ) -> None: """Init DeviceListener.""" self.hass = hass - self.entry = entry + self.device_manager = device_manager + self.device_ids = device_ids def update_device(self, device: TuyaDevice) -> None: """Update device status.""" - if device.id in self.hass.data[DOMAIN][self.entry.entry_id][TUYA_HA_DEVICES]: + if device.id in self.device_ids: _LOGGER.debug( "_update-->%s;->>%s", self, @@ -171,33 +178,14 @@ class DeviceListener(TuyaDeviceListener): def add_device(self, device: TuyaDevice) -> None: """Add device added listener.""" - device_add = False - - if device.category in itertools.chain( - *self.hass.data[DOMAIN][self.entry.entry_id][TUYA_HA_TUYA_MAP].values() - ): - ha_tuya_map = self.hass.data[DOMAIN][self.entry.entry_id][TUYA_HA_TUYA_MAP] + if device.category in TUYA_SUPPORTED_PRODUCT_CATEGORIES: + # Ensure the device isn't present stale self.hass.add_job(self.async_remove_device, device.id) - for domain, tuya_list in ha_tuya_map.items(): - if device.category in tuya_list: - device_add = True - _LOGGER.debug( - "Add device category->%s; domain-> %s", - device.category, - domain, - ) - self.hass.data[DOMAIN][self.entry.entry_id][TUYA_HA_DEVICES].add( - device.id - ) - dispatcher_send( - self.hass, TUYA_DISCOVERY_NEW.format(domain), [device.id] - ) + self.device_ids.add(device.id) + dispatcher_send(self.hass, TUYA_DISCOVERY_NEW, [device.id]) - if device_add: - device_manager = self.hass.data[DOMAIN][self.entry.entry_id][ - TUYA_DEVICE_MANAGER - ] + device_manager = self.device_manager device_manager.mq.stop() tuya_mq = TuyaOpenMQ(device_manager.api) tuya_mq.start() @@ -207,18 +195,16 @@ class DeviceListener(TuyaDeviceListener): def remove_device(self, device_id: str) -> None: """Add device removed listener.""" - _LOGGER.debug("tuya remove device:%s", device_id) self.hass.add_job(self.async_remove_device, device_id) @callback def async_remove_device(self, device_id: str) -> None: """Remove device from Home Assistant.""" + _LOGGER.debug("Remove device: %s", device_id) device_registry_object = device_registry.async_get(self.hass) device_entry = device_registry_object.async_get_device( identifiers={(DOMAIN, device_id)} ) if device_entry is not None: device_registry_object.async_remove_device(device_entry.id) - self.hass.data[DOMAIN][self.entry.entry_id][TUYA_HA_DEVICES].discard( - device_id - ) + self.device_ids.discard(device_id) diff --git a/homeassistant/components/tuya/climate.py b/homeassistant/components/tuya/climate.py index 810e8ad8aab..ac309464054 100644 --- a/homeassistant/components/tuya/climate.py +++ b/homeassistant/components/tuya/climate.py @@ -8,7 +8,7 @@ from typing import Any from tuya_iot import TuyaDevice, TuyaDeviceManager -from homeassistant.components.climate import DOMAIN as DEVICE_DOMAIN, ClimateEntity +from homeassistant.components.climate import ClimateEntity from homeassistant.components.climate.const import ( HVAC_MODE_AUTO, HVAC_MODE_COOL, @@ -25,17 +25,11 @@ from homeassistant.config_entries import ConfigEntry from homeassistant.const import TEMP_CELSIUS, TEMP_FAHRENHEIT from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.dispatcher import async_dispatcher_connect -from homeassistant.helpers.entity import Entity from homeassistant.helpers.entity_platform import AddEntitiesCallback +from . import HomeAssistantTuyaData from .base import TuyaHaEntity -from .const import ( - DOMAIN, - TUYA_DEVICE_MANAGER, - TUYA_DISCOVERY_NEW, - TUYA_HA_DEVICES, - TUYA_HA_TUYA_MAP, -) +from .const import DOMAIN, TUYA_DISCOVERY_NEW _LOGGER = logging.getLogger(__name__) @@ -88,50 +82,25 @@ TUYA_SUPPORT_TYPE = { async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: - """Set up tuya climate dynamically through tuya discovery.""" - _LOGGER.debug("climate init") - - hass.data[DOMAIN][entry.entry_id][TUYA_HA_TUYA_MAP][ - DEVICE_DOMAIN - ] = TUYA_SUPPORT_TYPE + """Set up Tuya climate dynamically through Tuya discovery.""" + hass_data: HomeAssistantTuyaData = hass.data[DOMAIN][entry.entry_id] @callback - def async_discover_device(dev_ids: list[str]) -> None: - """Discover and add a discovered tuya climate.""" - _LOGGER.debug("climate add-> %s", dev_ids) - if not dev_ids: - return - entities = _setup_entities(hass, entry, dev_ids) + def async_discover_device(device_ids: list[str]) -> None: + """Discover and add a discovered Tuya climate.""" + entities: list[TuyaHaClimate] = [] + for device_id in device_ids: + device = hass_data.device_manager.device_map[device_id] + if device and device.category in TUYA_SUPPORT_TYPE: + entities.append(TuyaHaClimate(device, hass_data.device_manager)) async_add_entities(entities) + async_discover_device([*hass_data.device_manager.device_map]) + entry.async_on_unload( - async_dispatcher_connect( - hass, TUYA_DISCOVERY_NEW.format(DEVICE_DOMAIN), async_discover_device - ) + async_dispatcher_connect(hass, TUYA_DISCOVERY_NEW, async_discover_device) ) - device_manager = hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER] - device_ids = [] - for (device_id, device) in device_manager.device_map.items(): - if device.category in TUYA_SUPPORT_TYPE: - device_ids.append(device_id) - async_discover_device(device_ids) - - -def _setup_entities( - hass: HomeAssistant, entry: ConfigEntry, device_ids: list[str] -) -> list[Entity]: - """Set up Tuya Climate.""" - device_manager = hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER] - entities: list[Entity] = [] - for device_id in device_ids: - device = device_manager.device_map[device_id] - if device is None: - continue - entities.append(TuyaHaClimate(device, device_manager)) - hass.data[DOMAIN][entry.entry_id][TUYA_HA_DEVICES].add(device_id) - return entities - class TuyaHaClimate(TuyaHaEntity, ClimateEntity): """Tuya Switch Device.""" diff --git a/homeassistant/components/tuya/const.py b/homeassistant/components/tuya/const.py index e4ecfbd4bf3..6d92383a871 100644 --- a/homeassistant/components/tuya/const.py +++ b/homeassistant/components/tuya/const.py @@ -15,12 +15,8 @@ CONF_PASSWORD = "password" CONF_COUNTRY_CODE = "country_code" CONF_APP_TYPE = "tuya_app_type" -TUYA_DISCOVERY_NEW = "tuya_discovery_new_{}" -TUYA_DEVICE_MANAGER = "tuya_device_manager" -TUYA_HOME_MANAGER = "tuya_home_manager" -TUYA_MQTT_LISTENER = "tuya_mqtt_listener" -TUYA_HA_TUYA_MAP = "tuya_ha_tuya_map" -TUYA_HA_DEVICES = "tuya_ha_devices" +TUYA_DISCOVERY_NEW = "tuya_discovery_new" +TUYA_HA_SIGNAL_UPDATE_ENTITY = "tuya_entry_update" TUYA_RESPONSE_CODE = "code" TUYA_RESPONSE_RESULT = "result" @@ -28,7 +24,29 @@ TUYA_RESPONSE_MSG = "msg" TUYA_RESPONSE_SUCCESS = "success" TUYA_RESPONSE_PLATFROM_URL = "platform_url" -TUYA_HA_SIGNAL_UPDATE_ENTITY = "tuya_entry_update" +TUYA_SUPPORTED_PRODUCT_CATEGORIES = ( + "bh", # Smart Kettle + "cwysj", # Pet Water Feeder + "cz", # Socket + "dc", # Light string + "dd", # Light strip + "dj", # Light + "dlq", # Breaker + "fs", # Fan + "fs", # Fan + "fwl", # Ambient light + "jsq", # Humidifier's light + "kg", # Switch + "kj", # Air Purifier + "kj", # Air Purifier + "kt", # Air conditioner + "pc", # Power Strip + "qn", # Heater + "wk", # Thermostat + "xdd", # Ceiling Light + "xxj", # Diffuser + "xxj", # Diffuser's light +) TUYA_SMART_APP = "tuyaSmart" SMARTLIFE_APP = "smartlife" diff --git a/homeassistant/components/tuya/fan.py b/homeassistant/components/tuya/fan.py index 15a8e553a10..88aa1ab53b9 100644 --- a/homeassistant/components/tuya/fan.py +++ b/homeassistant/components/tuya/fan.py @@ -10,7 +10,6 @@ from tuya_iot import TuyaDevice, TuyaDeviceManager from homeassistant.components.fan import ( DIRECTION_FORWARD, DIRECTION_REVERSE, - DOMAIN as DEVICE_DOMAIN, SUPPORT_DIRECTION, SUPPORT_OSCILLATE, SUPPORT_PRESET_MODE, @@ -26,14 +25,9 @@ from homeassistant.util.percentage import ( percentage_to_ordered_list_item, ) +from . import HomeAssistantTuyaData from .base import TuyaHaEntity -from .const import ( - DOMAIN, - TUYA_DEVICE_MANAGER, - TUYA_DISCOVERY_NEW, - TUYA_HA_DEVICES, - TUYA_HA_TUYA_MAP, -) +from .const import DOMAIN, TUYA_DISCOVERY_NEW _LOGGER = logging.getLogger(__name__) @@ -61,49 +55,24 @@ async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback ): """Set up tuya fan dynamically through tuya discovery.""" - _LOGGER.debug("fan init") - - hass.data[DOMAIN][entry.entry_id][TUYA_HA_TUYA_MAP][ - DEVICE_DOMAIN - ] = TUYA_SUPPORT_TYPE + hass_data: HomeAssistantTuyaData = hass.data[DOMAIN][entry.entry_id] @callback - def async_discover_device(dev_ids: list[str]) -> None: + def async_discover_device(device_ids: list[str]) -> None: """Discover and add a discovered tuya fan.""" - _LOGGER.debug("fan add-> %s", dev_ids) - if not dev_ids: - return - entities = _setup_entities(hass, entry, dev_ids) + entities: list[TuyaHaFan] = [] + for device_id in device_ids: + device = hass_data.device_manager.device_map[device_id] + if device and device.category in TUYA_SUPPORT_TYPE: + entities.append(TuyaHaFan(device, hass_data.device_manager)) async_add_entities(entities) + async_discover_device([*hass_data.device_manager.device_map]) + entry.async_on_unload( - async_dispatcher_connect( - hass, TUYA_DISCOVERY_NEW.format(DEVICE_DOMAIN), async_discover_device - ) + async_dispatcher_connect(hass, TUYA_DISCOVERY_NEW, async_discover_device) ) - device_manager = hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER] - device_ids = [] - for (device_id, device) in device_manager.device_map.items(): - if device.category in TUYA_SUPPORT_TYPE: - device_ids.append(device_id) - async_discover_device(device_ids) - - -def _setup_entities( - hass: HomeAssistant, entry: ConfigEntry, device_ids: list[str] -) -> list[TuyaHaFan]: - """Set up Tuya Fan.""" - device_manager = hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER] - entities = [] - for device_id in device_ids: - device = device_manager.device_map[device_id] - if device is None: - continue - entities.append(TuyaHaFan(device, device_manager)) - hass.data[DOMAIN][entry.entry_id][TUYA_HA_DEVICES].add(device_id) - return entities - class TuyaHaFan(TuyaHaEntity, FanEntity): """Tuya Fan Device.""" diff --git a/homeassistant/components/tuya/light.py b/homeassistant/components/tuya/light.py index f2136da431f..1970ba4c312 100644 --- a/homeassistant/components/tuya/light.py +++ b/homeassistant/components/tuya/light.py @@ -15,7 +15,6 @@ from homeassistant.components.light import ( COLOR_MODE_COLOR_TEMP, COLOR_MODE_HS, COLOR_MODE_ONOFF, - DOMAIN as DEVICE_DOMAIN, LightEntity, ) from homeassistant.config_entries import ConfigEntry @@ -23,14 +22,9 @@ from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity_platform import AddEntitiesCallback +from . import HomeAssistantTuyaData from .base import TuyaHaEntity -from .const import ( - DOMAIN, - TUYA_DEVICE_MANAGER, - TUYA_DISCOVERY_NEW, - TUYA_HA_DEVICES, - TUYA_HA_TUYA_MAP, -) +from .const import DOMAIN, TUYA_DISCOVERY_NEW _LOGGER = logging.getLogger(__name__) @@ -84,54 +78,24 @@ async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: """Set up tuya light dynamically through tuya discovery.""" - _LOGGER.debug("light init") - - hass.data[DOMAIN][entry.entry_id][TUYA_HA_TUYA_MAP][ - DEVICE_DOMAIN - ] = TUYA_SUPPORT_TYPE + hass_data: HomeAssistantTuyaData = hass.data[DOMAIN][entry.entry_id] @callback - def async_discover_device(dev_ids: list[str]): + def async_discover_device(device_ids: list[str]): """Discover and add a discovered tuya light.""" - _LOGGER.debug("light add-> %s", dev_ids) - if not dev_ids: - return - entities = _setup_entities(hass, entry, dev_ids) + entities: list[TuyaHaLight] = [] + for device_id in device_ids: + device = hass_data.device_manager.device_map[device_id] + if device and device.category in TUYA_SUPPORT_TYPE: + entities.append(TuyaHaLight(device, hass_data.device_manager)) async_add_entities(entities) + async_discover_device([*hass_data.device_manager.device_map]) + entry.async_on_unload( - async_dispatcher_connect( - hass, TUYA_DISCOVERY_NEW.format(DEVICE_DOMAIN), async_discover_device - ) + async_dispatcher_connect(hass, TUYA_DISCOVERY_NEW, async_discover_device) ) - device_manager = hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER] - device_ids = [] - for (device_id, device) in device_manager.device_map.items(): - if device.category in TUYA_SUPPORT_TYPE: - device_ids.append(device_id) - async_discover_device(device_ids) - - -def _setup_entities( - hass, entry: ConfigEntry, device_ids: list[str] -) -> list[TuyaHaLight]: - """Set up Tuya Light device.""" - device_manager = hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER] - entities = [] - for device_id in device_ids: - device = device_manager.device_map[device_id] - if device is None: - continue - - tuya_ha_light = TuyaHaLight(device, device_manager) - entities.append(tuya_ha_light) - hass.data[DOMAIN][entry.entry_id][TUYA_HA_DEVICES].add( - tuya_ha_light.tuya_device.id - ) - - return entities - class TuyaHaLight(TuyaHaEntity, LightEntity): """Tuya light device.""" diff --git a/homeassistant/components/tuya/scene.py b/homeassistant/components/tuya/scene.py index c90c6798b9b..f008d5a5b7b 100644 --- a/homeassistant/components/tuya/scene.py +++ b/homeassistant/components/tuya/scene.py @@ -1,7 +1,6 @@ """Support for Tuya scenes.""" from __future__ import annotations -import logging from typing import Any from tuya_iot import TuyaHomeManager, TuyaScene @@ -9,55 +8,48 @@ from tuya_iot import TuyaHomeManager, TuyaScene from homeassistant.components.scene import Scene from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant +from homeassistant.helpers.entity import DeviceInfo from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN, TUYA_HOME_MANAGER - -_LOGGER = logging.getLogger(__name__) +from . import HomeAssistantTuyaData +from .const import DOMAIN async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: - """Set up tuya scenes.""" - home_manager = hass.data[DOMAIN][entry.entry_id][TUYA_HOME_MANAGER] - scenes = await hass.async_add_executor_job(home_manager.query_scenes) - async_add_entities(TuyaHAScene(home_manager, scene) for scene in scenes) + """Set up Tuya scenes.""" + hass_data: HomeAssistantTuyaData = hass.data[DOMAIN][entry.entry_id] + scenes = await hass.async_add_executor_job(hass_data.home_manager.query_scenes) + async_add_entities(TuyaHAScene(hass_data.home_manager, scene) for scene in scenes) class TuyaHAScene(Scene): """Tuya Scene Remote.""" + _should_poll = False + def __init__(self, home_manager: TuyaHomeManager, scene: TuyaScene) -> None: """Init Tuya Scene.""" super().__init__() + self._attr_unique_id = f"tys{scene.scene_id}" self.home_manager = home_manager self.scene = scene - @property - def should_poll(self) -> bool: - """Hass should not poll.""" - return False - - @property - def unique_id(self) -> str | None: - """Return a unique ID.""" - return f"tys{self.scene.scene_id}" - @property def name(self) -> str | None: """Return Tuya scene name.""" return self.scene.name @property - def device_info(self): + def device_info(self) -> DeviceInfo: """Return a device description for device registry.""" - return { - "identifiers": {(DOMAIN, f"{self.unique_id}")}, - "manufacturer": "tuya", - "name": self.scene.name, - "model": "Tuya Scene", - } + return DeviceInfo( + identifiers={(DOMAIN, f"{self.unique_id}")}, + manufacturer="tuya", + name=self.scene.name, + model="Tuya Scene", + ) @property def available(self) -> bool: diff --git a/homeassistant/components/tuya/switch.py b/homeassistant/components/tuya/switch.py index 5bafbe1b7f6..af723a3d4f4 100644 --- a/homeassistant/components/tuya/switch.py +++ b/homeassistant/components/tuya/switch.py @@ -6,21 +6,15 @@ from typing import Any from tuya_iot import TuyaDevice, TuyaDeviceManager -from homeassistant.components.switch import DOMAIN as DEVICE_DOMAIN, SwitchEntity +from homeassistant.components.switch import SwitchEntity from homeassistant.config_entries import ConfigEntry -from homeassistant.core import HomeAssistant +from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.dispatcher import async_dispatcher_connect -from homeassistant.helpers.entity import Entity from homeassistant.helpers.entity_platform import AddEntitiesCallback +from . import HomeAssistantTuyaData from .base import TuyaHaEntity -from .const import ( - DOMAIN, - TUYA_DEVICE_MANAGER, - TUYA_DISCOVERY_NEW, - TUYA_HA_DEVICES, - TUYA_HA_TUYA_MAP, -) +from .const import DOMAIN, TUYA_DISCOVERY_NEW _LOGGER = logging.getLogger(__name__) @@ -62,45 +56,36 @@ async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: """Set up tuya sensors dynamically through tuya discovery.""" - _LOGGER.debug("switch init") + hass_data: HomeAssistantTuyaData = hass.data[DOMAIN][entry.entry_id] - hass.data[DOMAIN][entry.entry_id][TUYA_HA_TUYA_MAP][ - DEVICE_DOMAIN - ] = TUYA_SUPPORT_TYPE - - async def async_discover_device(dev_ids): + @callback + def async_discover_device(device_ids: list[str]) -> None: """Discover and add a discovered tuya sensor.""" - _LOGGER.debug("switch add-> %s", dev_ids) - if not dev_ids: - return - entities = _setup_entities(hass, entry, dev_ids) - async_add_entities(entities) + async_add_entities( + _setup_entities(hass, entry, hass_data.device_manager, device_ids) + ) + + async_discover_device([*hass_data.device_manager.device_map]) entry.async_on_unload( - async_dispatcher_connect( - hass, TUYA_DISCOVERY_NEW.format(DEVICE_DOMAIN), async_discover_device - ) + async_dispatcher_connect(hass, TUYA_DISCOVERY_NEW, async_discover_device) ) - device_manager = hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER] - device_ids = [] - for (device_id, device) in device_manager.device_map.items(): - if device.category in TUYA_SUPPORT_TYPE: - device_ids.append(device_id) - await async_discover_device(device_ids) - -def _setup_entities(hass, entry: ConfigEntry, device_ids: list[str]) -> list[Entity]: +def _setup_entities( + hass: HomeAssistant, + entry: ConfigEntry, + device_manager: TuyaDeviceManager, + device_ids: list[str], +) -> list[TuyaHaSwitch]: """Set up Tuya Switch device.""" - device_manager = hass.data[DOMAIN][entry.entry_id][TUYA_DEVICE_MANAGER] - entities: list[Entity] = [] + entities: list[TuyaHaSwitch] = [] for device_id in device_ids: device = device_manager.device_map[device_id] - if device is None: + if device is None or device.category not in TUYA_SUPPORT_TYPE: continue for function in device.function: - tuya_ha_switch = None if device.category == "kj": if function in [ DPCODE_ANION, @@ -110,26 +95,26 @@ def _setup_entities(hass, entry: ConfigEntry, device_ids: list[str]) -> list[Ent DPCODE_UV, DPCODE_WET, ]: - tuya_ha_switch = TuyaHaSwitch(device, device_manager, function) - # Main device switch is handled by the Fan object + entities.append(TuyaHaSwitch(device, device_manager, function)) + elif device.category == "cwysj": - if function in [DPCODE_FRESET, DPCODE_UV, DPCODE_PRESET, DPCODE_WRESET]: - tuya_ha_switch = TuyaHaSwitch(device, device_manager, function) + if ( + function + in [ + DPCODE_FRESET, + DPCODE_UV, + DPCODE_PRESET, + DPCODE_WRESET, + ] + or function.startswith(DPCODE_SWITCH) + ): + entities.append(TuyaHaSwitch(device, device_manager, function)) - if function.startswith(DPCODE_SWITCH): - # Main device switch - tuya_ha_switch = TuyaHaSwitch(device, device_manager, function) - else: - if function.startswith(DPCODE_START): - tuya_ha_switch = TuyaHaSwitch(device, device_manager, function) - if function.startswith(DPCODE_SWITCH): - tuya_ha_switch = TuyaHaSwitch(device, device_manager, function) + elif function.startswith(DPCODE_START) or function.startswith( + DPCODE_SWITCH + ): + entities.append(TuyaHaSwitch(device, device_manager, function)) - if tuya_ha_switch is not None: - entities.append(tuya_ha_switch) - hass.data[DOMAIN][entry.entry_id][TUYA_HA_DEVICES].add( - tuya_ha_switch.tuya_device.id - ) return entities