Xiaomi MiIO Switch: Support for different device types (#9836)

* Support for different device types of MIIO switches: Xiaomi Smart WiFi Socket (called Plug), Xiaomi Smart Power Strip and Xiaomi Chuang Mi Plug V1.

* Line too long fixed.

* Trailing whitespace removed.

* Changes based on review.

* Line too long fixed.

* No blank lines allowed after function docstring fixed.

* The underlying library is called python-miio now. Imports and requirements updated.

* TODO comment removed. Travis complains about.

* Blank line removed.

* Code clean-up.

* Revert "Code clean-up."

This reverts commit 96b191c7a6d30df3b2f8a301491ced61c84a49e2.

* Unused platform constant removed.

homeassistant/components/switch/xiaomi_miio.py

@@ -18,7 +18,7 @@ from homeassistant.exceptions import PlatformNotReady
 _LOGGER = logging.getLogger(__name__)
 
 DEFAULT_NAME = 'Xiaomi Miio Switch'
-PLATFORM = 'xiaomi_miio'
+
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
     vol.Required(CONF_HOST): cv.string,
     vol.Required(CONF_TOKEN): vol.All(cv.string, vol.Length(min=32, max=32)),
@@ -38,7 +38,7 @@ SUCCESS = ['ok']
 @asyncio.coroutine
 def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
     """Set up the switch from config."""
-    from mirobo import Plug, DeviceException
+    from miio import Device, DeviceException
 
     host = config.get(CONF_HOST)
     name = config.get(CONF_NAME)
@@ -46,23 +46,51 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
 
     _LOGGER.info("Initializing with host %s (token %s...)", host, token[:5])
 
+    devices = []
     try:
-        plug = Plug(host, token)
+        plug = Device(host, token)
         device_info = plug.info()
         _LOGGER.info("%s %s %s initialized",
-                     device_info.raw['model'],
-                     device_info.raw['fw_ver'],
-                     device_info.raw['hw_ver'])
+                     device_info.model,
+                     device_info.firmware_version,
+                     device_info.hardware_version)
 
-        xiaomi_plug_switch = XiaomiPlugSwitch(name, plug, device_info)
+        if device_info.model in ['chuangmi.plug.v1']:
+            from miio import PlugV1
+            plug = PlugV1(host, token)
+
+            # The device has two switchable channels (mains and a USB port).
+            # A switch device per channel will be created.
+            for channel_usb in [True, False]:
+                device = ChuangMiPlugV1Switch(
+                    name, plug, device_info, channel_usb)
+                devices.append(device)
+
+        elif device_info.model in ['qmi.powerstrip.v1',
+                                   'zimi.powerstrip.v2']:
+            from miio import Strip
+            plug = Strip(host, token)
+            device = XiaomiPowerStripSwitch(name, plug, device_info)
+            devices.append(device)
+        elif device_info.model in ['chuangmi.plug.m1',
+                                   'chuangmi.plug.v2']:
+            from miio import Plug
+            plug = Plug(host, token)
+            device = XiaomiPlugGenericSwitch(name, plug, device_info)
+            devices.append(device)
+        else:
+            _LOGGER.error(
+                'Unsupported device found! Please create an issue at '
+                'https://github.com/rytilahti/python-miio/issues '
+                'and provide the following data: %s', device_info.model)
     except DeviceException:
         raise PlatformNotReady
 
-    async_add_devices([xiaomi_plug_switch], update_before_add=True)
+    async_add_devices(devices, update_before_add=True)
 
 
-class XiaomiPlugSwitch(SwitchDevice):
-    """Representation of a Xiaomi Plug."""
+class XiaomiPlugGenericSwitch(SwitchDevice):
+    """Representation of a Xiaomi Plug Generic."""
 
     def __init__(self, name, plug, device_info):
         """Initialize the plug switch."""
@@ -74,8 +102,7 @@ class XiaomiPlugSwitch(SwitchDevice):
         self._state = None
         self._state_attrs = {
             ATTR_TEMPERATURE: None,
-            ATTR_LOAD_POWER: None,
-            ATTR_MODEL: self._device_info.raw['model'],
+            ATTR_MODEL: self._device_info.model,
         }
         self._skip_update = False
 
@@ -112,7 +139,7 @@ class XiaomiPlugSwitch(SwitchDevice):
     @asyncio.coroutine
     def _try_command(self, mask_error, func, *args, **kwargs):
         """Call a plug command handling error messages."""
-        from mirobo import DeviceException
+        from miio import DeviceException
         try:
             result = yield from self.hass.async_add_job(
                 partial(func, *args, **kwargs))
@@ -147,7 +174,43 @@ class XiaomiPlugSwitch(SwitchDevice):
     @asyncio.coroutine
     def async_update(self):
         """Fetch state from the device."""
-        from mirobo import DeviceException
+        from miio import DeviceException
+
+        # On state change the device doesn't provide the new state immediately.
+        if self._skip_update:
+            self._skip_update = False
+            return
+
+        try:
+            state = yield from self.hass.async_add_job(self._plug.status)
+            _LOGGER.debug("Got new state: %s", state)
+
+            self._state = state.is_on
+            self._state_attrs.update({
+                ATTR_TEMPERATURE: state.temperature
+            })
+
+        except DeviceException as ex:
+            _LOGGER.error("Got exception while fetching the state: %s", ex)
+
+
+class XiaomiPowerStripSwitch(XiaomiPlugGenericSwitch, SwitchDevice):
+    """Representation of a Xiaomi Power Strip."""
+
+    def __init__(self, name, plug, device_info):
+        """Initialize the plug switch."""
+        XiaomiPlugGenericSwitch.__init__(self, name, plug, device_info)
+
+        self._state_attrs = {
+            ATTR_TEMPERATURE: None,
+            ATTR_LOAD_POWER: None,
+            ATTR_MODEL: self._device_info.model,
+        }
+
+    @asyncio.coroutine
+    def async_update(self):
+        """Fetch state from the device."""
+        from miio import DeviceException
 
         # On state change the device doesn't provide the new state immediately.
         if self._skip_update:
@@ -161,8 +224,69 @@ class XiaomiPlugSwitch(SwitchDevice):
             self._state = state.is_on
             self._state_attrs.update({
                 ATTR_TEMPERATURE: state.temperature,
-                ATTR_LOAD_POWER: state.load_power,
+                ATTR_LOAD_POWER: state.load_power
             })
 
         except DeviceException as ex:
             _LOGGER.error("Got exception while fetching the state: %s", ex)
+
+
+class ChuangMiPlugV1Switch(XiaomiPlugGenericSwitch, SwitchDevice):
+    """Representation of a Chuang Mi Plug V1."""
+
+    def __init__(self, name, plug, device_info, channel_usb):
+        """Initialize the plug switch."""
+        name = name + ' USB' if channel_usb else name
+
+        XiaomiPlugGenericSwitch.__init__(self, name, plug, device_info)
+        self._channel_usb = channel_usb
+
+    @asyncio.coroutine
+    def async_turn_on(self, **kwargs):
+        """Turn a channel on."""
+        if self._channel_usb:
+            result = yield from self._try_command(
+                "Turning the plug on failed.", self._plug.usb_on)
+        else:
+            result = yield from self._try_command(
+                "Turning the plug on failed.", self._plug.on)
+
+        if result:
+            self._state = True
+            self._skip_update = True
+
+    @asyncio.coroutine
+    def async_turn_off(self, **kwargs):
+        """Turn a channel off."""
+        if self._channel_usb:
+            result = yield from self._try_command(
+                "Turning the plug on failed.", self._plug.usb_off)
+        else:
+            result = yield from self._try_command(
+                "Turning the plug on failed.", self._plug.off)
+
+        if result:
+            self._state = False
+            self._skip_update = True
+
+    @asyncio.coroutine
+    def async_update(self):
+        """Fetch state from the device."""
+        from miio import DeviceException
+
+        # On state change the device doesn't provide the new state immediately.
+        if self._skip_update:
+            self._skip_update = False
+            return
+
+        try:
+            state = yield from self.hass.async_add_job(self._plug.status)
+            _LOGGER.debug("Got new state: %s", state)
+
+            if self._channel_usb:
+                self._state = state.usb_power
+            else:
+                self._state = state.is_on
+
+        except DeviceException as ex:
+            _LOGGER.error("Got exception while fetching the state: %s", ex)