Make config flows general based on entities

custom_components/localtuya/__init__.py

@@ -5,8 +5,9 @@ import voluptuous as vol
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
+from homeassistant.const import CONF_PLATFORM
 
-from .const import CONF_DEVICE_TYPE, DOMAIN, DEVICE_TYPE_POWER_OUTLET
+from .const import DOMAIN, PLATFORMS
 
 
 async def async_setup(hass: HomeAssistant, config: dict):
@@ -16,11 +17,10 @@ async def async_setup(hass: HomeAssistant, config: dict):
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
     """Set up LocalTuya integration from a config entry."""
-    if entry.data[CONF_DEVICE_TYPE] == DEVICE_TYPE_POWER_OUTLET:
+    for component in PLATFORMS:
         hass.async_create_task(
-            hass.config_entries.async_forward_entry_setup(entry, "switch")
+            hass.config_entries.async_forward_entry_setup(entry, component)
         )
-
     return True
 
 

custom_components/localtuya/config_flow.py

@@ -1,75 +1,81 @@
 """Config flow for LocalTuya integration integration."""
 import logging
+from importlib import import_module
 
 import voluptuous as vol
 
 from homeassistant import config_entries, core, exceptions
 from homeassistant.const import (
-    CONF_SWITCHES,
+    CONF_ENTITIES,
     CONF_ID,
     CONF_HOST,
     CONF_DEVICE_ID,
     CONF_NAME,
     CONF_FRIENDLY_NAME,
+    CONF_PLATFORM,
+    CONF_SWITCHES,
 )
-import homeassistant.helpers.config_validation as cv
 
 from . import pytuya
 from .const import (  # pylint: disable=unused-import
-    CONF_DEVICE_TYPE,
     CONF_LOCAL_KEY,
     CONF_PROTOCOL_VERSION,
-    CONF_CURRENT,
-    CONF_CURRENT_CONSUMPTION,
-    CONF_VOLTAGE,
     DOMAIN,
-    DEVICE_TYPE_POWER_OUTLET,
+    PLATFORMS,
 )
 
 _LOGGER = logging.getLogger(__name__)
 
-ADD_ANOTHER_SWITCH = "add_another_switch"
+PLATFORM_TO_ADD = "platform_to_add"
+NO_ADDITIONAL_PLATFORMS = "no_additional_platforms"
 
 USER_SCHEMA = vol.Schema(
     {
-        vol.Required(CONF_NAME): str,
+        vol.Required(CONF_NAME, default="bla"): str,
-        vol.Required(CONF_HOST): str,
+        vol.Required(CONF_HOST, default="10.0.10.110"): str,
-        vol.Required(CONF_DEVICE_ID): str,
+        vol.Required(CONF_DEVICE_ID, default="307001182462ab4de00b"): str,
-        vol.Required(CONF_LOCAL_KEY): str,
+        vol.Required(CONF_LOCAL_KEY, default="28ac0651f2d187df"): str,
         vol.Required(CONF_PROTOCOL_VERSION, default="3.3"): vol.In(["3.1", "3.3"]),
-        vol.Required(CONF_DEVICE_TYPE, default=DEVICE_TYPE_POWER_OUTLET): vol.In(
-            [DEVICE_TYPE_POWER_OUTLET]
-        ),
     }
 )
 
-POWER_OUTLET_SCHEMA = vol.Schema(
+PICK_ENTITY_SCHEMA = vol.Schema(
-    {
+    {vol.Required(PLATFORM_TO_ADD, default=PLATFORMS[0]): vol.In(PLATFORMS)}
-        vol.Required(CONF_ID, default=1): int,
-        vol.Required(CONF_NAME): str,
-        vol.Required(CONF_FRIENDLY_NAME): str,
-        vol.Required(CONF_CURRENT, default=18): int,
-        vol.Required(CONF_CURRENT_CONSUMPTION, default=19): int,
-        vol.Required(CONF_VOLTAGE, default=20): int,
-        vol.Required(ADD_ANOTHER_SWITCH, default=False): bool,
-    }
 )
 
 
+def platform_schema(dps, additional_fields):
+    """Generate input validation schema for a platform."""
+    dps_list = vol.In([f"{id} (value: {value})" for id, value in dps.items()])
+    return vol.Schema(
+        {
+            vol.Required(CONF_ID): dps_list,
+            vol.Required(CONF_FRIENDLY_NAME): str,
+        }
+    ).extend({conf: dps_list for conf in additional_fields})
+
+
+def strip_dps_values(user_input, fields):
+    """Remove values and keep only index for DPS config items."""
+    for field in [CONF_ID] + fields:
+        user_input[field] = user_input[field].split(" ")[0]
+    return user_input
+
+
 async def validate_input(hass: core.HomeAssistant, data):
     """Validate the user input allows us to connect."""
     pytuyadevice = pytuya.OutletDevice(
         data[CONF_DEVICE_ID], data[CONF_HOST], data[CONF_LOCAL_KEY]
     )
     pytuyadevice.set_version(float(data[CONF_PROTOCOL_VERSION]))
-    pytuyadevice.set_dpsUsed({"1": None})
+    pytuyadevice.set_dpsUsed({})
     try:
-        await hass.async_add_executor_job(pytuyadevice.status)
+        data = await hass.async_add_executor_job(pytuyadevice.status)
-    except ConnectionRefusedError:
+    except (ConnectionRefusedError, ConnectionResetError):
         raise CannotConnect
     except ValueError:
         raise InvalidAuth
-    return data
+    return data["dps"]
 
 
 class LocaltuyaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
@@ -81,7 +87,10 @@ class LocaltuyaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     def __init__(self):
         """Initialize a new LocaltuyaConfigFlow."""
         self.basic_info = None
-        self.switches = []
+        self.dps_data = None
+        self.platform = None
+        self.platform_dps_fields = None
+        self.entities = []
 
     async def async_step_user(self, user_input=None):
         """Handle the initial step."""
@@ -91,11 +100,9 @@ class LocaltuyaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             self._abort_if_unique_id_configured()
 
             try:
-                self.basic_info = await validate_input(self.hass, user_input)
+                self.basic_info = user_input
-                if self.basic_info[CONF_DEVICE_TYPE] == "Power Outlet":
+                self.dps_data = await validate_input(self.hass, user_input)
-                    return await self.async_step_power_outlet()
+                return await self.async_step_pick_entity_type()
-
-                return self.async_abort(reason="unsupported_device_type")
             except CannotConnect:
                 errors["base"] = "cannot_connect"
             except InvalidAuth:
@@ -108,29 +115,88 @@ class LocaltuyaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             step_id="user", data_schema=USER_SCHEMA, errors=errors
         )
 
-    async def async_step_power_outlet(self, user_input=None):
+    async def async_step_pick_entity_type(self, user_input=None):
-        """Handle the initial step."""
+        """Handle asking if user wants to add another entity."""
+        if user_input is not None:
+            if user_input.get(NO_ADDITIONAL_PLATFORMS):
+                config = {**self.basic_info, CONF_ENTITIES: self.entities}
+                return self.async_create_entry(title=config[CONF_NAME], data=config)
+
+            self._set_platform(user_input[PLATFORM_TO_ADD])
+            return await self.async_step_add_entity()
+
+        # Add a checkbox that allows bailing out from config flow iff at least one
+        # entity has been added
+        schema = PICK_ENTITY_SCHEMA
+        if self.platform is not None:
+            schema = schema.extend(
+                {vol.Required(NO_ADDITIONAL_PLATFORMS, default=True): bool}
+            )
+
+        return self.async_show_form(step_id="pick_entity_type", data_schema=schema)
+
+    async def async_step_add_entity(self, user_input=None):
+        """Handle adding a new entity."""
         errors = {}
         if user_input is not None:
             already_configured = any(
-                switch[CONF_ID] == user_input[CONF_ID] for switch in self.switches
+                switch[CONF_ID] == user_input[CONF_ID] for switch in self.entities
             )
             if not already_configured:
-                add_another_switch = user_input.pop(ADD_ANOTHER_SWITCH)
+                user_input[CONF_PLATFORM] = self.platform
-                self.switches.append(user_input)
+                self.entities.append(
-                if not add_another_switch:
+                    strip_dps_values(user_input, self.platform_dps_fields)
-                    config = {**self.basic_info, CONF_SWITCHES: self.switches}
+                )
-                    return self.async_create_entry(title=config[CONF_NAME], data=config)
+                return await self.async_step_pick_entity_type()
-            else:
+
-                errors["base"] = "switch_already_configured"
+            errors["base"] = "entity_already_configured"
 
         return self.async_show_form(
-            step_id="power_outlet",
+            step_id="add_entity",
-            data_schema=POWER_OUTLET_SCHEMA,
+            data_schema=platform_schema(self.dps_data, self.platform_dps_fields),
             errors=errors,
-            description_placeholders={"number": len(self.switches) + 1},
+            description_placeholders={"platform": self.platform},
         )
 
+    async def async_step_import(self, user_input):
+        """Handle import from YAML."""
+
+        def _convert_entity(conf):
+            converted = {
+                CONF_ID: conf[CONF_ID],
+                CONF_FRIENDLY_NAME: conf[CONF_FRIENDLY_NAME],
+                CONF_PLATFORM: self.platform,
+            }
+            for field in self.platform_dps_fields:
+                converted[str(field)] = conf[field]
+            return converted
+
+        await self.async_set_unique_id(user_input[CONF_DEVICE_ID])
+        self._set_platform(user_input[CONF_PLATFORM])
+
+        if len(user_input[CONF_SWITCHES]) > 0:
+            for switch_conf in user_input[CONF_SWITCHES].values():
+                self.entities.append(_convert_entity(switch_conf))
+        else:
+            self.entities.append(_convert_entity(user_input))
+
+        config = {
+            CONF_NAME: f"{user_input[CONF_DEVICE_ID]} (import from configuration.yaml)",
+            CONF_HOST: user_input[CONF_HOST],
+            CONF_DEVICE_ID: user_input[CONF_DEVICE_ID],
+            CONF_LOCAL_KEY: user_input[CONF_LOCAL_KEY],
+            CONF_PROTOCOL_VERSION: user_input[CONF_PROTOCOL_VERSION],
+            CONF_ENTITIES: self.entities,
+        }
+        self._abort_if_unique_id_configured(updates=config)
+        return self.async_create_entry(title=config[CONF_NAME], data=config)
+
+    def _set_platform(self, platform):
+        self.platform = platform
+        self.platform_dps_fields = import_module(
+            f"homeassistant.components.localtuya.{platform}"
+        ).DPS_FIELDS
+
 
 class CannotConnect(exceptions.HomeAssistantError):
     """Error to indicate we cannot connect."""

custom_components/localtuya/const.py

@@ -1,6 +1,9 @@
 """Constants for localtuya integration."""
 
-CONF_DEVICE_TYPE = "device_type"
+ATTR_CURRENT = 'current'
+ATTR_CURRENT_CONSUMPTION = 'current_consumption'
+ATTR_VOLTAGE = 'voltage'
+
 CONF_LOCAL_KEY = "local_key"
 CONF_PROTOCOL_VERSION = "protocol_version"
 CONF_CURRENT = "current"
@@ -9,4 +12,5 @@ CONF_VOLTAGE = "voltage"
 
 DOMAIN = "localtuya"
 
-DEVICE_TYPE_POWER_OUTLET = "Power Outlet"
+# Platforms in this list must support config flows
+PLATFORMS = ["switch"]

custom_components/localtuya/switch.py

@@ -28,161 +28,131 @@ import logging
 import voluptuous as vol
 
 from homeassistant.components.switch import SwitchEntity, PLATFORM_SCHEMA
-from homeassistant.const import (CONF_HOST, CONF_ID, CONF_SWITCHES, CONF_DEVICE_ID, CONF_FRIENDLY_NAME, CONF_ICON, CONF_NAME)
+from homeassistant.const import (
+    CONF_HOST,
+    CONF_ID,
+    CONF_ENTITIES,
+    CONF_SWITCHES,
+    CONF_DEVICE_ID,
+    CONF_FRIENDLY_NAME,
+    CONF_ICON,
+    CONF_NAME,
+    CONF_PLATFORM,
+)
+from homeassistant.config_entries import SOURCE_IMPORT
 import homeassistant.helpers.config_validation as cv
 from time import time, sleep
 from threading import Lock
 
 from . import pytuya
-from .const import CONF_LOCAL_KEY, CONF_PROTOCOL_VERSION, CONF_CURRENT, CONF_CURRENT_CONSUMPTION, CONF_VOLTAGE
+from .const import (
+    ATTR_CURRENT,
+    ATTR_CURRENT_CONSUMPTION,
+    ATTR_VOLTAGE,
+    CONF_LOCAL_KEY,
+    CONF_PROTOCOL_VERSION,
+    CONF_CURRENT,
+    CONF_CURRENT_CONSUMPTION,
+    CONF_VOLTAGE,
+    DOMAIN,
+)
 
 _LOGGER = logging.getLogger(__name__)
 
-REQUIREMENTS = ['pytuya==7.0.9']
+REQUIREMENTS = ["pytuya==7.0.9"]
 
-DEFAULT_ID = '1'
+DEFAULT_ID = "1"
 DEFAULT_PROTOCOL_VERSION = 3.3
 
-ATTR_CURRENT = 'current'
+SWITCH_SCHEMA = vol.Schema(
-ATTR_CURRENT_CONSUMPTION = 'current_consumption'
+    {
-ATTR_VOLTAGE = 'voltage'
+        vol.Optional(CONF_ID, default=DEFAULT_ID): cv.string,
+        vol.Required(CONF_NAME): cv.string,
+        vol.Required(CONF_FRIENDLY_NAME): cv.string,
+        vol.Optional(CONF_CURRENT, default="-1"): cv.string,
+        vol.Optional(CONF_CURRENT_CONSUMPTION, default="-1"): cv.string,
+        vol.Optional(CONF_VOLTAGE, default="-1"): cv.string,
+    }
+)
 
-SWITCH_SCHEMA = vol.Schema({
+PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
-    vol.Optional(CONF_ID, default=DEFAULT_ID): cv.string,
+    {
-    vol.Required(CONF_NAME): cv.string,
+        vol.Optional(CONF_ICON): cv.icon,
-    vol.Required(CONF_FRIENDLY_NAME): cv.string,
+        vol.Required(CONF_HOST): cv.string,
-    vol.Optional(CONF_CURRENT, default='-1'): cv.string,
+        vol.Required(CONF_DEVICE_ID): cv.string,
-    vol.Optional(CONF_CURRENT_CONSUMPTION, default='-1'): cv.string,
+        vol.Required(CONF_LOCAL_KEY): cv.string,
-    vol.Optional(CONF_VOLTAGE, default='-1'): cv.string,
+        vol.Required(CONF_NAME): cv.string,
-})
+        vol.Required(CONF_FRIENDLY_NAME): cv.string,
+        vol.Required(
+            CONF_PROTOCOL_VERSION, default=DEFAULT_PROTOCOL_VERSION
+        ): vol.Coerce(float),
+        vol.Optional(CONF_ID, default=DEFAULT_ID): cv.string,
+        vol.Optional(CONF_CURRENT, default="-1"): cv.string,
+        vol.Optional(CONF_CURRENT_CONSUMPTION, default="-1"): cv.string,
+        vol.Optional(CONF_VOLTAGE, default="-1"): cv.string,
+        vol.Optional(CONF_SWITCHES, default={}): vol.Schema({cv.slug: SWITCH_SCHEMA}),
+    }
+)
 
-PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
+DPS_FIELDS = [
-    vol.Optional(CONF_ICON): cv.icon,
+    vol.Required(CONF_CURRENT),
-    vol.Required(CONF_HOST): cv.string,
+    vol.Required(CONF_CURRENT_CONSUMPTION),
-    vol.Required(CONF_DEVICE_ID): cv.string,
+    vol.Required(CONF_VOLTAGE),
-    vol.Required(CONF_LOCAL_KEY): cv.string,
+]
-    vol.Required(CONF_NAME): cv.string,
-    vol.Required(CONF_FRIENDLY_NAME): cv.string,
-    vol.Required(CONF_PROTOCOL_VERSION, default=DEFAULT_PROTOCOL_VERSION): vol.Coerce(float),
-    vol.Optional(CONF_ID, default=DEFAULT_ID): cv.string,
-    vol.Optional(CONF_CURRENT, default='-1'): cv.string,
-    vol.Optional(CONF_CURRENT_CONSUMPTION, default='-1'): cv.string,
-    vol.Optional(CONF_VOLTAGE, default='-1'): cv.string,
-    vol.Optional(CONF_SWITCHES, default={}):
-        vol.Schema({cv.slug: SWITCH_SCHEMA}),
-})
 
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Setup a Tuya switch based on a config entry."""
+    switches_to_setup = [
+        entity
+        for entity in config_entry.data[CONF_ENTITIES]
+        if entity[CONF_PLATFORM] == "switch"
+    ]
+    if not switches_to_setup:
+        return
+
     switches = []
     pytuyadevice = pytuya.OutletDevice(
         config_entry.data[CONF_DEVICE_ID],
         config_entry.data[CONF_HOST],
-        config_entry.data[CONF_LOCAL_KEY])
+        config_entry.data[CONF_LOCAL_KEY],
+    )
     pytuyadevice.set_version(float(config_entry.data[CONF_PROTOCOL_VERSION]))
-    dps = {}
+    pytuyadevice.set_dpsUsed({})
-
-    for device_config in config_entry.data[CONF_SWITCHES]:
-        dps[str(device_config[CONF_ID])] = None
-        if device_config[CONF_CURRENT] != 0:
-            dps[str(device_config[CONF_CURRENT])] = None
-        if device_config[CONF_CURRENT_CONSUMPTION] != 0:
-            dps[str(device_config[CONF_CURRENT_CONSUMPTION])] = None
-        if device_config[CONF_VOLTAGE] != 0:
-            dps[str(device_config[CONF_VOLTAGE])] = None
 
+    for device_config in switches_to_setup:
         switches.append(
-                TuyaDevice(
+            TuyaDevice(
-                    TuyaCache(pytuyadevice),
+                TuyaCache(pytuyadevice),
-                    device_config[CONF_NAME],
+                device_config[CONF_FRIENDLY_NAME],
-                    device_config[CONF_FRIENDLY_NAME],
+                device_config[CONF_ID],
-                    None,  # Icon
+                device_config[CONF_CURRENT],
-                    str(device_config[CONF_ID]),
+                device_config[CONF_CURRENT_CONSUMPTION],
-                    str(device_config[CONF_CURRENT]),
+                device_config[CONF_VOLTAGE],
-                    str(device_config[CONF_CURRENT_CONSUMPTION]),
+            )
-                    str(device_config[CONF_VOLTAGE])
-                )
         )
 
-    pytuyadevice.set_dpsUsed(dps)
-
     async_add_entities(switches, True)
 
 
 def setup_platform(hass, config, add_devices, discovery_info=None):
     """Set up of the Tuya switch."""
-    from . import pytuya
+    config[CONF_PLATFORM] = "switch"
-
+    hass.async_create_task(
-    devices = config.get(CONF_SWITCHES)
+        hass.config_entries.flow.async_init(
-
+            DOMAIN, context={"source": SOURCE_IMPORT}, data=config
-    switches = []
-    pytuyadevice = pytuya.OutletDevice(config.get(CONF_DEVICE_ID), config.get(CONF_HOST), config.get(CONF_LOCAL_KEY))
-    pytuyadevice.set_version(float(config.get(CONF_PROTOCOL_VERSION)))
-
-    if len(devices) > 0:
-        for object_id, device_config in devices.items():
-            dps = {}
-            dps[device_config.get(CONF_ID)]=None
-            if device_config.get(CONF_CURRENT) != '-1':
-                dps[device_config.get(CONF_CURRENT)]=None
-            if device_config.get(CONF_CURRENT_CONSUMPTION) != '-1':
-                dps[device_config.get(CONF_CURRENT_CONSUMPTION)]=None
-            if device_config.get(CONF_VOLTAGE) != '-1':
-                dps[device_config.get(CONF_VOLTAGE)]=None
-            pytuyadevice.set_dpsUsed(dps)
-
-            outlet_device = TuyaCache(pytuyadevice)
-            switches.append(
-                    TuyaDevice(
-                        outlet_device,
-                        device_config.get(CONF_NAME),
-                        device_config.get(CONF_FRIENDLY_NAME, object_id),
-                        device_config.get(CONF_ICON),
-                        device_config.get(CONF_ID),
-                        device_config.get(CONF_CURRENT),
-                        device_config.get(CONF_CURRENT_CONSUMPTION),
-                        device_config.get(CONF_VOLTAGE)
-                    )
-            )
-
-            print('Setup localtuya subswitch [{}] with device ID [{}] '.format(device_config.get(CONF_FRIENDLY_NAME, object_id), device_config.get(CONF_ID)))
-            _LOGGER.info("Setup localtuya subswitch %s with device ID %s ", device_config.get(CONF_FRIENDLY_NAME, object_id), device_config.get(CONF_ID) )
-    else:
-        dps = {}
-        dps[config.get(CONF_ID)]=None
-        if config.get(CONF_CURRENT) != '-1':
-            dps[config.get(CONF_CURRENT)]=None
-        if config.get(CONF_CURRENT_CONSUMPTION) != '-1':
-            dps[config.get(CONF_CURRENT_CONSUMPTION)]=None
-        if config.get(CONF_VOLTAGE) != '-1':
-            dps[config.get(CONF_VOLTAGE)]=None
-        pytuyadevice.set_dpsUsed(dps)
-        outlet_device = TuyaCache(pytuyadevice)
-        switches.append(
-                TuyaDevice(
-                    outlet_device,
-                    config.get(CONF_NAME),
-                    config.get(CONF_FRIENDLY_NAME),
-                    config.get(CONF_ICON),
-                    config.get(CONF_ID),
-                    config.get(CONF_CURRENT),
-                    config.get(CONF_CURRENT_CONSUMPTION),
-                    config.get(CONF_VOLTAGE)
-                )
         )
+    )
 
-        print('Setup localtuya switch [{}] with device ID [{}] '.format(config.get(CONF_FRIENDLY_NAME), config.get(CONF_ID)))
+    return True
-        _LOGGER.info("Setup localtuya switch %s with device ID %s ", config.get(CONF_FRIENDLY_NAME), config.get(CONF_ID) )
 
-    add_devices(switches, True)
 
 class TuyaCache:
     """Cache wrapper for pytuya.OutletDevice"""
 
     def __init__(self, device):
         """Initialize the cache."""
-        self._cached_status = ''
+        self._cached_status = ""
         self._cached_status_time = 0
         self._device = device
         self._lock = Lock()
@@ -198,26 +168,37 @@ class TuyaCache:
                 status = self._device.status()
                 return status
             except Exception:
-                print('Failed to update status of device [{}]'.format(self._device.address))
+                print(
+                    "Failed to update status of device [{}]".format(
+                        self._device.address
+                    )
+                )
                 sleep(1.0)
-                if i+1 == 3:
+                if i + 1 == 3:
-                    _LOGGER.error("Failed to update status of device %s", self._device.address )
+                    _LOGGER.error(
-#                    return None
+                        "Failed to update status of device %s", self._device.address
+                    )
+                    #                    return None
                     raise ConnectionError("Failed to update status .")
 
     def set_status(self, state, switchid):
         """Change the Tuya switch status and clear the cache."""
-        self._cached_status = ''
+        self._cached_status = ""
         self._cached_status_time = 0
         for i in range(5):
             try:
                 return self._device.set_status(state, switchid)
             except Exception:
-                print('Failed to set status of device [{}]'.format(self._device.address))
+                print(
-                if i+1 == 3:
+                    "Failed to set status of device [{}]".format(self._device.address)
-                    _LOGGER.error("Failed to set status of device %s", self._device.address )
+                )
+                if i + 1 == 3:
+                    _LOGGER.error(
+                        "Failed to set status of device %s", self._device.address
+                    )
                     return
-#                    raise ConnectionError("Failed to set status.")
+
+    #                    raise ConnectionError("Failed to set status.")
 
     def status(self):
         """Get state of Tuya switch and cache the results."""
@@ -236,19 +217,30 @@ class TuyaCache:
 class TuyaDevice(SwitchEntity):
     """Representation of a Tuya switch."""
 
-    def __init__(self, device, name, friendly_name, icon, switchid, attr_current, attr_consumption, attr_voltage):
+    def __init__(
+        self,
+        device,
+        friendly_name,
+        switchid,
+        attr_current,
+        attr_consumption,
+        attr_voltage,
+    ):
         """Initialize the Tuya switch."""
         self._device = device
         self._name = friendly_name
         self._available = False
-        self._icon = icon
         self._switch_id = switchid
         self._attr_current = attr_current
         self._attr_consumption = attr_consumption
         self._attr_voltage = attr_voltage
         self._status = None
         self._state = None
-        print('Initialized tuya switch [{}] with switch status [{}] and state [{}]'.format(self._name, self._status, self._state))
+        print(
+            "Initialized tuya switch [{}] with switch status [{}] and state [{}]".format(
+                self._name, self._status, self._state
+            )
+        )
 
     @property
     def name(self):
@@ -274,22 +266,21 @@ class TuyaDevice(SwitchEntity):
     def device_state_attributes(self):
         attrs = {}
         try:
-            attrs[ATTR_CURRENT] = "{}".format(self._status['dps'][self._attr_current])
+            attrs[ATTR_CURRENT] = "{}".format(self._status["dps"][self._attr_current])
-            attrs[ATTR_CURRENT_CONSUMPTION] = "{}".format(self._status['dps'][self._attr_consumption]/10)
+            attrs[ATTR_CURRENT_CONSUMPTION] = "{}".format(
-            attrs[ATTR_VOLTAGE] = "{}".format(self._status['dps'][self._attr_voltage]/10)
+                self._status["dps"][self._attr_consumption] / 10
-#            print('attrs[ATTR_CURRENT]: [{}]'.format(attrs[ATTR_CURRENT]))
+            )
-#            print('attrs[ATTR_CURRENT_CONSUMPTION]: [{}]'.format(attrs[ATTR_CURRENT_CONSUMPTION]))
+            attrs[ATTR_VOLTAGE] = "{}".format(
-#            print('attrs[ATTR_VOLTAGE]: [{}]'.format(attrs[ATTR_VOLTAGE]))
+                self._status["dps"][self._attr_voltage] / 10
+            )
+        #            print('attrs[ATTR_CURRENT]: [{}]'.format(attrs[ATTR_CURRENT]))
+        #            print('attrs[ATTR_CURRENT_CONSUMPTION]: [{}]'.format(attrs[ATTR_CURRENT_CONSUMPTION]))
+        #            print('attrs[ATTR_VOLTAGE]: [{}]'.format(attrs[ATTR_VOLTAGE]))
 
         except KeyError:
             pass
         return attrs
 
-    @property
-    def icon(self):
-        """Return the icon."""
-        return self._icon
-
     def turn_on(self, **kwargs):
         """Turn Tuya switch on."""
         self._device.set_status(True, self._switch_id)
@@ -302,7 +293,7 @@ class TuyaDevice(SwitchEntity):
         """Get state of Tuya switch."""
         try:
             self._status = self._device.status()
-            self._state = self._status['dps'][self._switch_id]
+            self._state = self._status["dps"][self._switch_id]
         except Exception:
             self._available = False
         else:

custom_components/localtuya/translations/en.json

@@ -1,39 +1,43 @@
 {
     "config": {
         "abort": {
-            "already_configured": "Device has already been configured.",
+            "already_configured": "Device has already been configured."
-            "unsupported_device_type": "Unsupported device type!"
         },
         "error": {
-            "cannot_connect": "Cannot connect to device. Verify that address is correct.",
+            "cannot_connect": "Cannot connect to device. Verify that address is correct and try again.",
             "invalid_auth": "Failed to authenticate with device. Verify that device id and local key are correct.",
             "unknown": "An unknown error occurred. See log for details.",
-            "switch_already_configured": "Switch with this ID has already been configured."
+            "entity_already_configured": "Entity with this ID has already been configured."
         },
         "step": {
             "user": {
                 "title": "Add Tuya device",
-                "description": "Fill in the basic details and pick device type. The name entered here will be used to identify the integration itself (as seen in the `Integrations` page). You will name each sub-device in the following steps.",
+                "description": "Fill in the basic details and pick device type. The name entered here will be used to identify the integration itself (as seen in the `Integrations` page). You will add entities and give them names in the following steps.",
                 "data": {
                     "name": "Name",
                     "host": "Host",
                     "device_id": "Device ID",
                     "local_key": "Local key",
-                    "protocol_version": "Protocol Version",
+                    "protocol_version": "Protocol Version"
-                    "device_type": "Device type"
                 }
             },
-            "power_outlet": {
+            "pick_entity_type": {
-                "title": "Add subswitch",
+                "title": "Entity type selection",
-		"description": "You are about to add subswitch number `{number}`. If you want to add another, tick `Add another switch` before continuing.",
+                "description": "Please pick the type of entity you want to add.",
+                "data": {
+                    "platform_to_add": "Platform",
+                    "no_additional_platforms": "Do not add any more entities"
+                }
+            },
+            "add_entity": {
+                "title": "Add new entity",
+                "description": "Please fill out the details for an entity with type `{platform}`.",
                 "data": {
                     "id": "ID",
-                    "name": "Name",
                     "friendly_name": "Friendly name",
                     "current": "Current",
                     "current_consumption": "Current Consumption",
-                    "voltage": "Voltage",
+                    "voltage": "Voltage"
-                    "add_another_switch": "Add another switch"
                 }
             }
         }