From d2ed3a131f1db07fe4d664053a26ecbd11a296bb Mon Sep 17 00:00:00 2001
From: Conrad Juhl Andersen <cnrd@users.noreply.github.com>
Date: Wed, 17 May 2017 08:26:57 +0200
Subject: [PATCH] Add support for disabling tradfri groups (#7593)

* Add support for disabling tradfri groups

* Fixed styleguide problems

* Fix style problems

* Use default for groups when setting up in UI
---
 homeassistant/components/light/tradfri.py |  8 +++++---
 homeassistant/components/tradfri.py       | 24 ++++++++++++++++++-----
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/homeassistant/components/light/tradfri.py b/homeassistant/components/light/tradfri.py
index fdd02d73349..28c42c5699c 100644
--- a/homeassistant/components/light/tradfri.py
+++ b/homeassistant/components/light/tradfri.py
@@ -11,7 +11,7 @@ from homeassistant.components.light import (
     SUPPORT_COLOR_TEMP, SUPPORT_RGB_COLOR, Light)
 from homeassistant.components.light import \
     PLATFORM_SCHEMA as LIGHT_PLATFORM_SCHEMA
-from homeassistant.components.tradfri import KEY_GATEWAY
+from homeassistant.components.tradfri import KEY_GATEWAY, KEY_TRADFRI_GROUPS
 from homeassistant.util import color as color_util
 
 _LOGGER = logging.getLogger(__name__)
@@ -35,8 +35,10 @@ def setup_platform(hass, config, add_devices, discovery_info=None):
     lights = [dev for dev in devices if dev.has_light_control]
     add_devices(Tradfri(light) for light in lights)
 
-    groups = gateway.get_groups()
-    add_devices(TradfriGroup(group) for group in groups)
+    allow_tradfri_groups = hass.data[KEY_TRADFRI_GROUPS][gateway_id]
+    if allow_tradfri_groups:
+        groups = gateway.get_groups()
+        add_devices(TradfriGroup(group) for group in groups)
 
 
 class TradfriGroup(Light):
diff --git a/homeassistant/components/tradfri.py b/homeassistant/components/tradfri.py
index e1ef0f5fabd..cd83f81afd1 100644
--- a/homeassistant/components/tradfri.py
+++ b/homeassistant/components/tradfri.py
@@ -17,16 +17,22 @@ from homeassistant.const import CONF_HOST, CONF_API_KEY
 from homeassistant.loader import get_component
 from homeassistant.components.discovery import SERVICE_IKEA_TRADFRI
 
+REQUIREMENTS = ['pytradfri==1.1']
+
 DOMAIN = 'tradfri'
 CONFIG_FILE = 'tradfri.conf'
 KEY_CONFIG = 'tradfri_configuring'
 KEY_GATEWAY = 'tradfri_gateway'
-REQUIREMENTS = ['pytradfri==1.1']
+KEY_TRADFRI_GROUPS = 'tradfri_allow_tradfri_groups'
+CONF_ALLOW_TRADFRI_GROUPS = 'allow_tradfri_groups'
+DEFAULT_ALLOW_TRADFRI_GROUPS = True
 
 CONFIG_SCHEMA = vol.Schema({
     DOMAIN: vol.Schema({
         vol.Inclusive(CONF_HOST, 'gateway'): cv.string,
         vol.Inclusive(CONF_API_KEY, 'gateway'): cv.string,
+        vol.Optional(CONF_ALLOW_TRADFRI_GROUPS,
+                     default=DEFAULT_ALLOW_TRADFRI_GROUPS): cv.boolean,
     })
 }, extra=vol.ALLOW_EXTRA)
 
@@ -47,7 +53,8 @@ def request_configuration(hass, config, host):
     def configuration_callback(callback_data):
         """Handle the submitted configuration."""
         res = yield from _setup_gateway(hass, config, host,
-                                        callback_data.get('key'))
+                                        callback_data.get('key'),
+                                        DEFAULT_ALLOW_TRADFRI_GROUPS)
         if not res:
             hass.async_add_job(configurator.notify_errors, instance,
                                "Unable to connect.")
@@ -77,6 +84,7 @@ def async_setup(hass, config):
     conf = config.get(DOMAIN, {})
     host = conf.get(CONF_HOST)
     key = conf.get(CONF_API_KEY)
+    allow_tradfri_groups = conf.get(CONF_ALLOW_TRADFRI_GROUPS)
 
     @asyncio.coroutine
     def gateway_discovered(service, info):
@@ -85,7 +93,8 @@ def async_setup(hass, config):
         host = info['host']
 
         if host in keys:
-            yield from _setup_gateway(hass, config, host, keys[host]['key'])
+            yield from _setup_gateway(hass, config, host, keys[host]['key'],
+                                      allow_tradfri_groups)
         else:
             hass.async_add_job(request_configuration, hass, config, host)
 
@@ -94,11 +103,12 @@ def async_setup(hass, config):
     if host is None:
         return True
 
-    return (yield from _setup_gateway(hass, config, host, key))
+    return (yield from _setup_gateway(hass, config, host, key,
+                                      allow_tradfri_groups))
 
 
 @asyncio.coroutine
-def _setup_gateway(hass, hass_config, host, key):
+def _setup_gateway(hass, hass_config, host, key, allow_tradfri_groups):
     """Create a gateway."""
     from pytradfri import cli_api_factory, Gateway, RequestError, retry_timeout
 
@@ -112,6 +122,10 @@ def _setup_gateway(hass, hass_config, host, key):
     hass.data.setdefault(KEY_GATEWAY, {})
     gateways = hass.data[KEY_GATEWAY]
 
+    hass.data.setdefault(KEY_TRADFRI_GROUPS, {})
+    tradfri_groups = hass.data[KEY_TRADFRI_GROUPS]
+    tradfri_groups[gateway_id] = allow_tradfri_groups
+
     # Check if already set up
     if gateway_id in gateways:
         return True