Refactor async_setup_entry

custom_components/localtuya/binary_sensor.py

@@ -1,5 +1,6 @@
 """Platform to present any Tuya DP as a binary sensor."""
 import logging
+from functools import partial
 
 import voluptuous as vol
 
@@ -8,9 +9,9 @@ from homeassistant.components.binary_sensor import (
     DEVICE_CLASSES_SCHEMA,
     BinarySensorEntity,
 )
-from homeassistant.const import CONF_ID, CONF_DEVICE_CLASS
+from homeassistant.const import CONF_DEVICE_CLASS
 
-from .common import LocalTuyaEntity, prepare_setup_entities
+from .common import LocalTuyaEntity, async_setup_entry
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -27,27 +28,6 @@ def flow_schema(dps):
     }
 
 
-async def async_setup_entry(hass, config_entry, async_add_entities):
-    """Set up a Tuya sensor based on a config entry."""
-    tuyainterface, entities_to_setup = prepare_setup_entities(
-        hass, config_entry, DOMAIN
-    )
-    if not entities_to_setup:
-        return
-
-    sensors = []
-    for device_config in entities_to_setup:
-        sensors.append(
-            LocaltuyaBinarySensor(
-                tuyainterface,
-                config_entry,
-                device_config[CONF_ID],
-            )
-        )
-
-    async_add_entities(sensors)
-
-
 class LocaltuyaBinarySensor(LocalTuyaEntity, BinarySensorEntity):
     """Representation of a Tuya binary sensor."""
 
@@ -83,3 +63,6 @@ class LocaltuyaBinarySensor(LocalTuyaEntity, BinarySensorEntity):
             _LOGGER.warning(
                 "State for entity %s did not match state patterns", self.entity_id
             )
+
+
+async_setup_entry = partial(async_setup_entry, DOMAIN, LocaltuyaBinarySensor)

custom_components/localtuya/common.py

@@ -39,6 +39,33 @@ def prepare_setup_entities(hass, config_entry, platform):
     return tuyainterface, entities_to_setup
 
 
+async def async_setup_entry(
+    domain, entity_class, hass, config_entry, async_add_entities
+):
+    """Set up a Tuya platform based on a config entry.
+
+    This is a generic method and each platform should lock domain and
+    entity_class with functools.partial.
+    """
+    tuyainterface, entities_to_setup = prepare_setup_entities(
+        hass, config_entry, domain
+    )
+    if not entities_to_setup:
+        return
+
+    entities = []
+    for device_config in entities_to_setup:
+        entities.append(
+            entity_class(
+                tuyainterface,
+                config_entry,
+                device_config[CONF_ID],
+            )
+        )
+
+    async_add_entities(entities)
+
+
 def get_entity_config(config_entry, dps_id):
     """Return entity config for a given DPS id."""
     for entity in config_entry.data[CONF_ENTITIES]:

custom_components/localtuya/cover.py

@@ -1,5 +1,6 @@
 """Platform to locally control Tuya-based cover devices."""
 import logging
+from functools import partial
 from time import sleep
 
 import voluptuous as vol
@@ -13,7 +14,6 @@ from homeassistant.components.cover import (
     SUPPORT_SET_POSITION,
     ATTR_POSITION,
 )
-from homeassistant.const import CONF_ID
 
 from .const import (
     CONF_OPENCLOSE_CMDS,
@@ -22,7 +22,7 @@ from .const import (
     CONF_POSITIONING_MODE,
     CONF_SPAN_TIME,
 )
-from .common import LocalTuyaEntity, prepare_setup_entities
+from .common import LocalTuyaEntity, async_setup_entry
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -55,27 +55,6 @@ def flow_schema(dps):
     }
 
 
-async def async_setup_entry(hass, config_entry, async_add_entities):
-    """Set up a Tuya cover based on a config entry."""
-    tuyainterface, entities_to_setup = prepare_setup_entities(
-        hass, config_entry, DOMAIN
-    )
-    if not entities_to_setup:
-        return
-
-    covers = []
-    for device_config in entities_to_setup:
-        covers.append(
-            LocaltuyaCover(
-                tuyainterface,
-                config_entry,
-                device_config[CONF_ID],
-            )
-        )
-
-    async_add_entities(covers)
-
-
 class LocaltuyaCover(LocalTuyaEntity, CoverEntity):
     """Tuya cover device."""
 
@@ -184,3 +163,6 @@ class LocaltuyaCover(LocalTuyaEntity, CoverEntity):
             )
         else:
             self._current_cover_position = 50
+
+
+async_setup_entry = partial(async_setup_entry, DOMAIN, LocaltuyaCover)

custom_components/localtuya/fan.py

@@ -1,5 +1,6 @@
 """Platform to locally control Tuya-based fan devices."""
 import logging
+from functools import partial
 
 from homeassistant.components.fan import (
     FanEntity,
@@ -11,9 +12,8 @@ from homeassistant.components.fan import (
     SUPPORT_SET_SPEED,
     SUPPORT_OSCILLATE,
 )
-from homeassistant.const import CONF_ID
 
-from .common import LocalTuyaEntity, prepare_setup_entities
+from .common import LocalTuyaEntity, async_setup_entry
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -23,28 +23,6 @@ def flow_schema(dps):
     return {}
 
 
-async def async_setup_entry(hass, config_entry, async_add_entities):
-    """Set up a Tuya fan based on a config entry."""
-    tuyainterface, entities_to_setup = prepare_setup_entities(
-        hass, config_entry, DOMAIN
-    )
-    if not entities_to_setup:
-        return
-
-    fans = []
-
-    for device_config in entities_to_setup:
-        fans.append(
-            LocaltuyaFan(
-                tuyainterface,
-                config_entry,
-                device_config[CONF_ID],
-            )
-        )
-
-    async_add_entities(fans)
-
-
 class LocaltuyaFan(LocalTuyaEntity, FanEntity):
     """Representation of a Tuya fan."""
 
@@ -130,3 +108,6 @@ class LocaltuyaFan(LocalTuyaEntity, FanEntity):
         elif self._status["dps"]["2"] == "3":
             self._speed = SPEED_HIGH
         self._oscillating = self._status["dps"]["8"]
+
+
+async_setup_entry = partial(async_setup_entry, DOMAIN, LocaltuyaFan)

custom_components/localtuya/light.py

@@ -1,7 +1,7 @@
 """Platform to locally control Tuya-based light devices."""
 import logging
+from functools import partial
 
-from homeassistant.const import CONF_ID
 from homeassistant.components.light import (
     LightEntity,
     DOMAIN,
@@ -12,7 +12,7 @@ from homeassistant.components.light import (
     SUPPORT_COLOR,
 )
 
-from .common import LocalTuyaEntity, prepare_setup_entities
+from .common import LocalTuyaEntity, async_setup_entry
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -32,27 +32,6 @@ def flow_schema(dps):
     return {}
 
 
-async def async_setup_entry(hass, config_entry, async_add_entities):
-    """Set up a Tuya light based on a config entry."""
-    tuyainterface, entities_to_setup = prepare_setup_entities(
-        hass, config_entry, DOMAIN
-    )
-    if not entities_to_setup:
-        return
-
-    lights = []
-    for device_config in entities_to_setup:
-        lights.append(
-            LocaltuyaLight(
-                tuyainterface,
-                config_entry,
-                device_config[CONF_ID],
-            )
-        )
-
-    async_add_entities(lights)
-
-
 class LocaltuyaLight(LocalTuyaEntity, LightEntity):
     """Representation of a Tuya light."""
 
@@ -140,3 +119,6 @@ class LocaltuyaLight(LocalTuyaEntity, LightEntity):
         self._brightness = brightness
 
         self._color_temp = self.dps(DPS_INDEX_COLOURTEMP)
+
+
+async_setup_entry = partial(async_setup_entry, DOMAIN, LocaltuyaLight)

custom_components/localtuya/sensor.py

@@ -1,18 +1,18 @@
 """Platform to present any Tuya DP as a sensor."""
 import logging
+from functools import partial
 
 import voluptuous as vol
 
 from homeassistant.components.sensor import DOMAIN, DEVICE_CLASSES
 from homeassistant.const import (
-    CONF_ID,
     CONF_DEVICE_CLASS,
     CONF_UNIT_OF_MEASUREMENT,
     STATE_UNKNOWN,
 )
 
 from .const import CONF_SCALING
-from .common import LocalTuyaEntity, prepare_setup_entities
+from .common import LocalTuyaEntity, async_setup_entry
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -31,27 +31,6 @@ def flow_schema(dps):
     }
 
 
-async def async_setup_entry(hass, config_entry, async_add_entities):
-    """Set up a Tuya sensor based on a config entry."""
-    tuyainterface, entities_to_setup = prepare_setup_entities(
-        hass, config_entry, DOMAIN
-    )
-    if not entities_to_setup:
-        return
-
-    sensors = []
-    for device_config in entities_to_setup:
-        sensors.append(
-            LocaltuyaSensor(
-                tuyainterface,
-                config_entry,
-                device_config[CONF_ID],
-            )
-        )
-
-    async_add_entities(sensors)
-
-
 class LocaltuyaSensor(LocalTuyaEntity):
     """Representation of a Tuya sensor."""
 
@@ -88,3 +67,6 @@ class LocaltuyaSensor(LocalTuyaEntity):
         if scale_factor is not None:
             state = round(state * scale_factor, DEFAULT_PRECISION)
         self._state = state
+
+
+async_setup_entry = partial(async_setup_entry, DOMAIN, LocaltuyaSensor)

custom_components/localtuya/switch.py

@@ -1,5 +1,6 @@
 """Platform to locally control Tuya-based switch devices."""
 import logging
+from functools import partial
 
 import voluptuous as vol
 
@@ -7,7 +8,6 @@ from homeassistant.components.switch import (
     SwitchEntity,
     DOMAIN,
 )
-from homeassistant.const import CONF_ID
 
 from .const import (
     ATTR_CURRENT,
@@ -17,7 +17,7 @@ from .const import (
     CONF_CURRENT_CONSUMPTION,
     CONF_VOLTAGE,
 )
-from .common import LocalTuyaEntity, prepare_setup_entities
+from .common import LocalTuyaEntity, async_setup_entry
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -31,27 +31,6 @@ def flow_schema(dps):
     }
 
 
-async def async_setup_entry(hass, config_entry, async_add_entities):
-    """Set up a Tuya switch based on a config entry."""
-    tuyainterface, entities_to_setup = prepare_setup_entities(
-        hass, config_entry, DOMAIN
-    )
-    if not entities_to_setup:
-        return
-
-    switches = []
-    for device_config in entities_to_setup:
-        switches.append(
-            LocaltuyaSwitch(
-                tuyainterface,
-                config_entry,
-                device_config[CONF_ID],
-            )
-        )
-
-    async_add_entities(switches)
-
-
 class LocaltuyaSwitch(LocalTuyaEntity, SwitchEntity):
     """Representation of a Tuya switch."""
 
@@ -97,3 +76,6 @@ class LocaltuyaSwitch(LocalTuyaEntity, SwitchEntity):
     def status_updated(self):
         """Device status was updated."""
         self._state = self.dps(self._dps_id)
+
+
+async_setup_entry = partial(async_setup_entry, DOMAIN, LocaltuyaSwitch)