Merge branch 'pytuya_refactoring' into physical_device

custom_components/localtuya/switch.py

@@ -22,51 +22,48 @@ switch:
       sw02:
         name: usb_plug
         friendly_name: USB Plug
-        id: 7  
+        id: 7
 """
 import logging
-import voluptuous as vol
-
-from homeassistant.components.switch import SwitchEntity, PLATFORM_SCHEMA
-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
+import voluptuous as vol
+
+from homeassistant.components.switch import (
+    SwitchEntity,
+    PLATFORM_SCHEMA,
+)
+from homeassistant.const import (
+    CONF_ID,
+    CONF_SWITCHES,
+    CONF_FRIENDLY_NAME,
+    CONF_NAME,
+)
+import homeassistant.helpers.config_validation as cv
+
+from . import BASE_PLATFORM_SCHEMA, prepare_setup_entities, import_from_yaml
 from .const import (
     ATTR_CURRENT,
     ATTR_CURRENT_CONSUMPTION,
     ATTR_VOLTAGE,
-    CONF_LOCAL_KEY,
-    CONF_PROTOCOL_VERSION,
     CONF_CURRENT,
     CONF_CURRENT_CONSUMPTION,
     CONF_VOLTAGE,
-    DOMAIN,
 )
+from .pytuya import TuyaDevice
 
 _LOGGER = logging.getLogger(__name__)
-REQUIREMENTS = ['pytuya>=8.0.0']
+
+PLATFORM = "switch"
 
 DEFAULT_ID = "1"
-DEFAULT_PROTOCOL_VERSION = 3.3
 
+# TODO: This will eventully merge with flow_schema
 SWITCH_SCHEMA = vol.Schema(
     {
         vol.Optional(CONF_ID, default=DEFAULT_ID): cv.string,
-        vol.Required(CONF_NAME): cv.string,
+        vol.Optional(CONF_NAME): cv.string,  # Deprecated: not used
         vol.Required(CONF_FRIENDLY_NAME): cv.string,
         vol.Optional(CONF_CURRENT, default="-1"): cv.string,
         vol.Optional(CONF_CURRENT_CONSUMPTION, default="-1"): cv.string,
@@ -74,18 +71,8 @@ SWITCH_SCHEMA = vol.Schema(
     }
 )
 
-PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
+PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(BASE_PLATFORM_SCHEMA).extend(
     {
-        vol.Optional(CONF_ICON): cv.icon,
-        vol.Required(CONF_HOST): cv.string,
-        vol.Required(CONF_DEVICE_ID): cv.string,
-        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,
@@ -94,44 +81,28 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
 )
 
 
-DPS_FIELDS = [
-    vol.Optional(CONF_CURRENT),
-    vol.Optional(CONF_CURRENT_CONSUMPTION),
-    vol.Optional(CONF_VOLTAGE),
-]
+def flow_schema(dps):
+    """Return schema used in config flow."""
+    return {
+        vol.Optional(CONF_CURRENT): vol.In(dps),
+        vol.Optional(CONF_CURRENT_CONSUMPTION): vol.In(dps),
+        vol.Optional(CONF_VOLTAGE): vol.In(dps),
+    }
 
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Setup a Tuya switch based on a config entry."""
-
-    print('config_entry: [{}] '.format(config_entry.data))
-
-
-    switches_to_setup = [
-        entity
-        for entity in config_entry.data[CONF_ENTITIES]
-        if entity[CONF_PLATFORM] == "switch"
-    ]
-    if not switches_to_setup:
+    device, entities_to_setup = prepare_setup_entities(
+        config_entry, PLATFORM
+    )
+    if not entities_to_setup:
         return
 
     switches = []
-    pytuyadevice = pytuya.PytuyaDevice(
-        config_entry.data[CONF_DEVICE_ID],
-        config_entry.data[CONF_HOST],
-        config_entry.data[CONF_LOCAL_KEY],
-        config_entry.data[CONF_FRIENDLY_NAME],
-        config_entry.data[CONF_NAME],
-    )
-    pytuyadevice.set_version(float(config_entry.data[CONF_PROTOCOL_VERSION]))
-    pytuyadevice.set_dpsUsed({})
-
-    print('switches_to_setup: [{}] '.format(switches_to_setup))
-
-    for device_config in switches_to_setup:
+    for device_config in entities_to_setup:
         switches.append(
             LocaltuyaSwitch(
-                TuyaCache(pytuyadevice),
+                TuyaCache(device),
                 device_config[CONF_FRIENDLY_NAME],
                 device_config[CONF_ID],
                 device_config.get(CONF_CURRENT),
@@ -145,18 +116,11 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
 
 def setup_platform(hass, config, add_devices, discovery_info=None):
     """Set up of the Tuya switch."""
-    config[CONF_PLATFORM] = "switch"
-    hass.async_create_task(
-        hass.config_entries.flow.async_init(
-            DOMAIN, context={"source": SOURCE_IMPORT}, data=config
-        )
-    )
-
-    return True
+    return import_from_yaml(hass, config, PLATFORM)
 
 
 class TuyaCache:
-    """Cache wrapper for pytuya.PytuyaDevice"""
+    """Cache wrapper for pytuya.TuyaDevice"""
 
     def __init__(self, device):
         """Initialize the cache."""
@@ -189,13 +153,13 @@ class TuyaCache:
                     #                    return None
                     raise ConnectionError("Failed to update status .")
 
-    def set_dps(self, state, switchid):
+    def set_dps(self, state, dps_index):
         """Change the Tuya switch status and clear the cache."""
         self._cached_status = ""
         self._cached_status_time = 0
         for i in range(5):
             try:
-                return self._device.set_dps(state, switchid)
+                return self._device.set_dps(state, dps_index)
             except Exception:
                 print(
                     "Failed to set status of device [{}]".format(self._device.address)