Merge pull request #1 from leeyuentuen/branch2

Branch2
SheaSmith · Feb 9, 2022 · c84d10b · c84d10b
2 parents 0d7a213 + 6898579
commit c84d10b
Show file tree

Hide file tree

Showing 7 changed files with 688 additions and 177 deletions.
diff --git a/custom_components/localtuya/__init__.py b/custom_components/localtuya/__init__.py
@@ -6,10 +6,12 @@
 localtuya:
   - host: 192.168.1.x
     device_id: xxxxx
-    local_key: xxxxx
+    local_key: xxxxx # Optional
     friendly_name: Tuya Device
     protocol_version: "3.3"
-    entities:
+    is_gateway: false # Optional
+    parent_gateway: xxxxx # Optional
+    entities: # Optional
       - platform: binary_sensor
         friendly_name: Plug Status
         id: 1
@@ -90,9 +92,9 @@
 from homeassistant.helpers.event import async_track_time_interval
 from homeassistant.helpers.reload import async_integration_yaml_config
 
-from .common import TuyaDevice, async_config_entry_by_device_id
+from .common import TuyaDevice, TuyaGatewayDevice, TuyaSubDevice, async_config_entry_by_device_id
 from .config_flow import config_schema
-from .const import CONF_PRODUCT_KEY, DATA_DISCOVERY, DOMAIN, TUYA_DEVICE
+from .const import CONF_PRODUCT_KEY, CONF_IS_GATEWAY, CONF_PARENT_GATEWAY, DATA_DISCOVERY, DOMAIN, TUYA_DEVICE
 from .discovery import TuyaDiscovery
 
 _LOGGER = logging.getLogger(__name__)
@@ -260,30 +262,35 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
     """Set up LocalTuya integration from a config entry."""
     unsub_listener = entry.add_update_listener(update_listener)
 
-    device = TuyaDevice(hass, entry.data)
+    if entry.data.get(CONF_IS_GATEWAY):
+        device = TuyaGatewayDevice(hass, entry.data)
+    elif not entry.data.get(CONF_IS_GATEWAY) and entry.data.get(CONF_PARENT_GATEWAY):
+        device = TuyaSubDevice(hass, entry.data)
+    else:
+        device = TuyaDevice(hass, entry.data)
 
     hass.data[DOMAIN][entry.entry_id] = {
         UNSUB_LISTENER: unsub_listener,
         TUYA_DEVICE: device,
     }
 
-    async def setup_entities():
-        platforms = set(entity[CONF_PLATFORM] for entity in entry.data[CONF_ENTITIES])
-        await asyncio.gather(
-            *[
-                hass.config_entries.async_forward_entry_setup(entry, platform)
-                for platform in platforms
-            ]
-        )
-        device.async_connect()
-
-    await async_remove_orphan_entities(hass, entry)
-
-    hass.async_create_task(setup_entities())
+    if not entry.data.get(CONF_IS_GATEWAY):
+        async def setup_entities():
+            platforms = set(
+                entity[CONF_PLATFORM] for entity in entry.data[CONF_ENTITIES]
+            )
+            await asyncio.gather(
+                *[
+                    hass.config_entries.async_forward_entry_setup(entry, platform)
+                    for platform in platforms
+                ]
+            )
+            device.async_connect()
 
+        await async_remove_orphan_entities(hass, entry)
+        hass.async_create_task(setup_entities())
     return True
 
-
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry):
     """Unload a config entry."""
     unload_ok = all(
@@ -302,9 +309,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry):
     if unload_ok:
         hass.data[DOMAIN].pop(entry.entry_id)
 
-    return True
-
-
+    return True    
 async def update_listener(hass, config_entry):
     """Update listener."""
     await hass.config_entries.async_reload(config_entry.entry_id)