Improve Wemo setup speed

homeassistant/components/binary_sensor/wemo.py

@@ -18,7 +18,7 @@
 _LOGGER = logging.getLogger(__name__)
 
 
-def setup_platform(hass, config, add_entities_callback, discovery_info=None):
+def setup_platform(hass, config, add_entities, discovery_info=None):
     """Register discovered WeMo binary sensors."""
     from pywemo import discovery
 
@@ -34,7 +34,7 @@ def setup_platform(hass, config, add_entities_callback, discovery_info=None):
             raise PlatformNotReady
 
         if device:
-            add_entities_callback([WemoBinarySensor(hass, device)])
+            add_entities([WemoBinarySensor(hass, device)])
 
 
 class WemoBinarySensor(BinarySensorDevice):

homeassistant/components/switch/wemo.py

@@ -35,7 +35,7 @@
 WEMO_STANDBY = 8
 
 
-def setup_platform(hass, config, add_entities_callback, discovery_info=None):
+def setup_platform(hass, config, add_entities, discovery_info=None):
     """Set up discovered WeMo switches."""
     from pywemo import discovery
 
@@ -51,7 +51,7 @@ def setup_platform(hass, config, add_entities_callback, discovery_info=None):
             raise PlatformNotReady
 
         if device:
-            add_entities_callback([WemoSwitch(device)])
+            add_entities([WemoSwitch(device)])
 
 
 class WemoSwitch(SwitchDevice):

homeassistant/components/wemo.py

@@ -10,10 +10,11 @@
 import voluptuous as vol
 
 from homeassistant.components.discovery import SERVICE_WEMO
-from homeassistant.helpers import discovery
 from homeassistant.helpers import config_validation as cv
+from homeassistant.helpers import discovery
 
-from homeassistant.const import EVENT_HOMEASSISTANT_STOP
+from homeassistant.const import (
+    EVENT_HOMEASSISTANT_START, EVENT_HOMEASSISTANT_STOP)
 
 REQUIREMENTS = ['pywemo==0.4.34']
 
@@ -77,19 +78,37 @@ def setup(hass, config):
     """Set up for WeMo devices."""
     import pywemo
 
+    # Keep track of WeMo devices
+    devices = []
+
+    # Keep track of WeMo device subscriptions for push updates
     global SUBSCRIPTION_REGISTRY
     SUBSCRIPTION_REGISTRY = pywemo.SubscriptionRegistry()
     SUBSCRIPTION_REGISTRY.start()
 
     def stop_wemo(event):
         """Shutdown Wemo subscriptions and subscription thread on exit."""
-        _LOGGER.debug("Shutting down subscriptions.")
+        _LOGGER.debug("Shutting down WeMo event subscriptions")
         SUBSCRIPTION_REGISTRY.stop()
 
     hass.bus.listen_once(EVENT_HOMEASSISTANT_STOP, stop_wemo)
 
+    def setup_url_for_device(device):
+        """Determine setup.xml url for given device."""
+        return 'http://{}:{}/setup.xml'.format(device.host, device.port)
+
+    def setup_url_for_address(host, port):
+        """Determine setup.xml url for given host and port pair."""
+        if not port:
+            port = pywemo.ouimeaux_device.probe_wemo(host)
+
+        if not port:
+            return None
+
+        return 'http://{}:{}/setup.xml'.format(host, port)
+
     def discovery_dispatch(service, discovery_info):
-        """Dispatcher for WeMo discovery events."""
+        """Dispatcher for incoming WeMo discovery events."""
         # name, model, location, mac
         model_name = discovery_info.get('model_name')
         serial = discovery_info.get('serial')
@@ -99,69 +118,65 @@ def discovery_dispatch(service, discovery_info):
             _LOGGER.debug('Ignoring known device %s %s',
                           service, discovery_info)
             return
-        _LOGGER.debug('Discovered unique device %s', serial)
+
+        _LOGGER.debug('Discovered unique WeMo device: %s', serial)
         KNOWN_DEVICES.append(serial)
 
         component = WEMO_MODEL_DISPATCH.get(model_name, 'switch')
 
-        discovery.load_platform(hass, component, DOMAIN, discovery_info,
-                                config)
+        discovery.load_platform(hass, component, DOMAIN,
+                                discovery_info, config)
 
     discovery.listen(hass, SERVICE_WEMO, discovery_dispatch)
 
-    def setup_url_for_device(device):
-        """Determine setup.xml url for given device."""
-        return 'http://{}:{}/setup.xml'.format(device.host, device.port)
+    def discover_wemo_devices(now):
+        """Run discovery for WeMo devices."""
+        _LOGGER.debug("Beginning WeMo device discovery...")
+        _LOGGER.debug("Adding statically configured WeMo devices...")
+        for host, port in config.get(DOMAIN, {}).get(CONF_STATIC, []):
+            url = setup_url_for_address(host, port)
+
+            if not url:
+                _LOGGER.error(
+                    'Unable to get description url for WeMo at: %s',
+                    '{}:{}'.format(host, port) if port else host)
+                continue
+
+            try:
+                device = pywemo.discovery.device_from_description(url, None)
+            except (requests.exceptions.ConnectionError,
+                    requests.exceptions.Timeout) as err:
+                _LOGGER.error('Unable to access WeMo at %s (%s)', url, err)
+                continue
 
-    def setup_url_for_address(host, port):
-        """Determine setup.xml url for given host and port pair."""
-        if not port:
-            port = pywemo.ouimeaux_device.probe_wemo(host)
+            if not [d[1] for d in devices
+                    if d[1].serialnumber == device.serialnumber]:
+                devices.append((url, device))
 
-        if not port:
-            return None
+        if config.get(DOMAIN, {}).get(CONF_DISCOVERY):
+            _LOGGER.debug("Scanning network for WeMo devices...")
+            for device in pywemo.discover_devices():
+                if not [d[1] for d in devices
+                        if d[1].serialnumber == device.serialnumber]:
+                    devices.append((setup_url_for_device(device),
+                                    device))
 
-        return 'http://{}:{}/setup.xml'.format(host, port)
+        for url, device in devices:
+            _LOGGER.debug('Adding WeMo device at %s:%i',
+                          device.host, device.port)
 
-    devices = []
+            discovery_info = {
+                'model_name': device.model_name,
+                'serial': device.serialnumber,
+                'mac_address': device.mac,
+                'ssdp_description': url,
+            }
 
-    _LOGGER.debug("Scanning statically configured WeMo devices...")
-    for host, port in config.get(DOMAIN, {}).get(CONF_STATIC, []):
-        url = setup_url_for_address(host, port)
-
-        if not url:
-            _LOGGER.error(
-                'Unable to get description url for %s',
-                '{}:{}'.format(host, port) if port else host)
-            continue
-
-        try:
-            device = pywemo.discovery.device_from_description(url, None)
-        except (requests.exceptions.ConnectionError,
-                requests.exceptions.Timeout) as err:
-            _LOGGER.error('Unable to access %s (%s)', url, err)
-            continue
-
-        if not [d[1] for d in devices
-                if d[1].serialnumber == device.serialnumber]:
-            devices.append((url, device))
-
-    if config.get(DOMAIN, {}).get(CONF_DISCOVERY):
-        _LOGGER.debug("Scanning for WeMo devices...")
-        for device in pywemo.discover_devices():
-            if not [d[1] for d in devices
-                    if d[1].serialnumber == device.serialnumber]:
-                devices.append((setup_url_for_device(device), device))
+            discovery.discover(hass, SERVICE_WEMO, discovery_info)
 
-    for url, device in devices:
-        _LOGGER.debug('Adding WeMo device at %s:%i', device.host, device.port)
+        _LOGGER.debug("WeMo device discovery has finished")
 
-        discovery_info = {
-            'model_name': device.model_name,
-            'serial': device.serialnumber,
-            'mac_address': device.mac,
-            'ssdp_description': url,
-        }
+    hass.bus.listen_once(EVENT_HOMEASSISTANT_START,
+                         discover_wemo_devices)
 
-        discovery.discover(hass, SERVICE_WEMO, discovery_info)
     return True