From ad91e4b417889d6da27c9934c47bb7d95af8742b Mon Sep 17 00:00:00 2001
From: jan iversen <jancasacondor@gmail.com>
Date: Tue, 9 Nov 2021 15:18:13 +0100
Subject: [PATCH] Simplify setup of tradfri entities (#59343)

* Simplify detection of devices.
---
 homeassistant/components/tradfri/cover.py  |  6 +++---
 homeassistant/components/tradfri/fan.py    | 10 +++++-----
 homeassistant/components/tradfri/light.py  | 10 +++++-----
 homeassistant/components/tradfri/sensor.py | 18 +++++++++---------
 homeassistant/components/tradfri/switch.py |  8 +++-----
 5 files changed, 25 insertions(+), 27 deletions(-)

diff --git a/homeassistant/components/tradfri/cover.py b/homeassistant/components/tradfri/cover.py
index 074b6ddd726..554650f9005 100644
--- a/homeassistant/components/tradfri/cover.py
+++ b/homeassistant/components/tradfri/cover.py
@@ -26,9 +26,9 @@ async def async_setup_entry(
     api = tradfri_data[KEY_API]
     devices = tradfri_data[DEVICES]
 
-    covers = [dev for dev in devices if dev.has_blind_control]
-    if covers:
-        async_add_entities(TradfriCover(cover, api, gateway_id) for cover in covers)
+    async_add_entities(
+        TradfriCover(dev, api, gateway_id) for dev in devices if dev.has_blind_control
+    )
 
 
 class TradfriCover(TradfriBaseDevice, CoverEntity):
diff --git a/homeassistant/components/tradfri/fan.py b/homeassistant/components/tradfri/fan.py
index d1b5de8b7cc..845d5e6d9c3 100644
--- a/homeassistant/components/tradfri/fan.py
+++ b/homeassistant/components/tradfri/fan.py
@@ -31,11 +31,11 @@ async def async_setup_entry(
     api = tradfri_data[KEY_API]
     devices = tradfri_data[DEVICES]
 
-    purifiers = [dev for dev in devices if dev.has_air_purifier_control]
-    if purifiers:
-        async_add_entities(
-            TradfriAirPurifierFan(purifier, api, gateway_id) for purifier in purifiers
-        )
+    async_add_entities(
+        TradfriAirPurifierFan(dev, api, gateway_id)
+        for dev in devices
+        if dev.has_air_purifier_control
+    )
 
 
 def _from_percentage(percentage: int) -> int:
diff --git a/homeassistant/components/tradfri/light.py b/homeassistant/components/tradfri/light.py
index 5a2da96152f..ca078a37e81 100644
--- a/homeassistant/components/tradfri/light.py
+++ b/homeassistant/components/tradfri/light.py
@@ -49,12 +49,12 @@ async def async_setup_entry(
     api = tradfri_data[KEY_API]
     devices = tradfri_data[DEVICES]
 
-    lights = [dev for dev in devices if dev.has_light_control]
-    if lights:
-        async_add_entities(TradfriLight(light, api, gateway_id) for light in lights)
-
+    entities: list[TradfriBaseClass] = [
+        TradfriLight(dev, api, gateway_id) for dev in devices if dev.has_light_control
+    ]
     if config_entry.data[CONF_IMPORT_GROUPS] and (groups := tradfri_data[GROUPS]):
-        async_add_entities(TradfriGroup(group, api, gateway_id) for group in groups)
+        entities.extend([TradfriGroup(group, api, gateway_id) for group in groups])
+    async_add_entities(entities)
 
 
 class TradfriGroup(TradfriBaseClass, LightEntity):
diff --git a/homeassistant/components/tradfri/sensor.py b/homeassistant/components/tradfri/sensor.py
index c6214773b97..27d2690ce34 100644
--- a/homeassistant/components/tradfri/sensor.py
+++ b/homeassistant/components/tradfri/sensor.py
@@ -27,17 +27,17 @@ async def async_setup_entry(
     api = tradfri_data[KEY_API]
     devices = tradfri_data[DEVICES]
 
-    sensors = (
-        dev
+    async_add_entities(
+        TradfriSensor(dev, api, gateway_id)
         for dev in devices
-        if not dev.has_light_control
-        and not dev.has_socket_control
-        and not dev.has_blind_control
-        and not dev.has_signal_repeater_control
-        and not dev.has_air_purifier_control
+        if (
+            not dev.has_light_control
+            and not dev.has_socket_control
+            and not dev.has_blind_control
+            and not dev.has_signal_repeater_control
+            and not dev.has_air_purifier_control
+        )
     )
-    if sensors:
-        async_add_entities(TradfriSensor(sensor, api, gateway_id) for sensor in sensors)
 
 
 class TradfriSensor(TradfriBaseDevice, SensorEntity):
diff --git a/homeassistant/components/tradfri/switch.py b/homeassistant/components/tradfri/switch.py
index d308e78c3d1..72d311fe555 100644
--- a/homeassistant/components/tradfri/switch.py
+++ b/homeassistant/components/tradfri/switch.py
@@ -26,11 +26,9 @@ async def async_setup_entry(
     api = tradfri_data[KEY_API]
     devices = tradfri_data[DEVICES]
 
-    switches = [dev for dev in devices if dev.has_socket_control]
-    if switches:
-        async_add_entities(
-            TradfriSwitch(switch, api, gateway_id) for switch in switches
-        )
+    async_add_entities(
+        TradfriSwitch(dev, api, gateway_id) for dev in devices if dev.has_socket_control
+    )
 
 
 class TradfriSwitch(TradfriBaseDevice, SwitchEntity):