Fix zha color probe

homeassistant/components/light/zha.py

@@ -31,42 +31,23 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
         return
 
     endpoint = discovery_info['endpoint']
-    try:
-        discovery_info['color_capabilities'] \
-            = yield from endpoint.light_color['color_capabilities']
-    except (AttributeError, KeyError):
-        pass
-
-    if discovery_info.get('color_capabilities') is None:
-        # ZCL Version 4 devices don't support the color_capabilities attribute.
-        # In this version XY support is mandatory, but we need to probe to
-        # determine if the device supports color temperature.
-        discovery_info['color_capabilities'] = CAPABILITIES_COLOR_XY
-        result = yield from safe_read(
-            endpoint.light_color, ['color_temperature'])
-        if result.get('color_temperature') is not UNSUPPORTED_ATTRIBUTE:
-            discovery_info['color_capabilities'] |= CAPABILITIES_COLOR_TEMP
+    if hasattr(endpoint, 'light_color'):
+        caps = yield from zha.safe_read(
+            endpoint.light_color, ['color_capabilities'])
+        discovery_info['color_capabilities'] = caps.get('color_capabilities')
+        if discovery_info['color_capabilities'] is None:
+            # ZCL Version 4 devices don't support the color_capabilities
+            # attribute. In this version XY support is mandatory, but we need
+            # to probe to determine if the device supports color temperature.
+            discovery_info['color_capabilities'] = CAPABILITIES_COLOR_XY
+            result = yield from zha.safe_read(
+                endpoint.light_color, ['color_temperature'])
+            if result.get('color_temperature') is not UNSUPPORTED_ATTRIBUTE:
+                discovery_info['color_capabilities'] |= CAPABILITIES_COLOR_TEMP
 
     async_add_devices([Light(**discovery_info)], update_before_add=True)
 
 
-@asyncio.coroutine
-def safe_read(cluster, attributes):
-    """Swallow all exceptions from network read.
-
-    If we throw during initialization, setup fails. Rather have an
-    entity that exists, but is in a maybe wrong state, than no entity.
-    """
-    try:
-        result, _ = yield from cluster.read_attributes(
-            attributes,
-            allow_cache=False,
-        )
-        return result
-    except Exception:  # pylint: disable=broad-except
-        return {}
-
-
 class Light(zha.Entity, light.Light):
     """Representation of a ZHA or ZLL light."""
 
@@ -174,23 +155,23 @@ def supported_features(self):
     @asyncio.coroutine
     def async_update(self):
         """Retrieve latest state."""
-        result = yield from safe_read(self._endpoint.on_off, ['on_off'])
+        result = yield from zha.safe_read(self._endpoint.on_off, ['on_off'])
         self._state = result.get('on_off', self._state)
 
         if self._supported_features & light.SUPPORT_BRIGHTNESS:
-            result = yield from safe_read(self._endpoint.level,
-                                          ['current_level'])
+            result = yield from zha.safe_read(self._endpoint.level,
+                                              ['current_level'])
             self._brightness = result.get('current_level', self._brightness)
 
         if self._supported_features & light.SUPPORT_COLOR_TEMP:
-            result = yield from safe_read(self._endpoint.light_color,
-                                          ['color_temperature'])
+            result = yield from zha.safe_read(self._endpoint.light_color,
+                                              ['color_temperature'])
             self._color_temp = result.get('color_temperature',
                                           self._color_temp)
 
         if self._supported_features & light.SUPPORT_XY_COLOR:
-            result = yield from safe_read(self._endpoint.light_color,
-                                          ['current_x', 'current_y'])
+            result = yield from zha.safe_read(self._endpoint.light_color,
+                                              ['current_x', 'current_y'])
             if 'current_x' in result and 'current_y' in result:
                 self._xy_color = (result['current_x'], result['current_y'])
 

homeassistant/components/zha/__init__.py

@@ -315,3 +315,21 @@ def get_discovery_info(hass, discovery_info):
     all_discovery_info = hass.data.get(DISCOVERY_KEY, {})
     discovery_info = all_discovery_info.get(discovery_key, None)
     return discovery_info
+
+
+@asyncio.coroutine
+def safe_read(cluster, attributes):
+    """Swallow all exceptions from network read.
+
+    If we throw during initialization, setup fails. Rather have an entity that
+    exists, but is in a maybe wrong state, than no entity. This method should
+    probably only be used during initialization.
+    """
+    try:
+        result, _ = yield from cluster.read_attributes(
+            attributes,
+            allow_cache=False,
+        )
+        return result
+    except Exception:  # pylint: disable=broad-except
+        return {}