Clarified get_registry code.

home-assistant · Mar 24, 2019 · 896bced · 896bced
1 parent cb4c178
commit 896bced
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 18 deletions.
diff --git a/homeassistant/helpers/area_registry.py b/homeassistant/helpers/area_registry.py
@@ -135,20 +135,23 @@ def _data_to_save(self) -> dict:
 @bind_hass
 async def async_get_registry(hass: HomeAssistantType) -> AreaRegistry:
     """Return area registry instance."""
-    reg = hass.data.get(DATA_REGISTRY)
+    reg_or_evt = hass.data.get(DATA_REGISTRY)
 
-    if not reg:
+    if not reg_or_evt:
         evt = hass.data[DATA_REGISTRY] = Event()
 
         reg = AreaRegistry(hass)
         await reg.async_load()
 
         hass.data[DATA_REGISTRY] = reg
         evt.set()
-    elif isinstance(reg, Event):
+        return reg
+
+    if isinstance(reg_or_evt, Event):
         # This shouldn't be possible in production (has been tested)
         # It's here soley for testing purposes.
-        await reg.wait()
-        reg = hass.data.get(DATA_REGISTRY)
+        evt = cast(Event, reg_or_evt)
+        await evt.wait()
+        return hass.data.get(DATA_REGISTRY)
 
-    return cast(AreaRegistry, reg)
+    return cast(AreaRegistry, reg_or_evt)
diff --git a/homeassistant/helpers/device_registry.py b/homeassistant/helpers/device_registry.py
@@ -277,23 +277,26 @@ def async_clear_area_id(self, area_id: str) -> None:
 @bind_hass
 async def async_get_registry(hass: HomeAssistantType) -> DeviceRegistry:
     """Return device registry instance."""
-    reg = hass.data.get(DATA_REGISTRY)
+    reg_or_evt = hass.data.get(DATA_REGISTRY)
 
-    if not reg:
+    if not reg_or_evt:
         evt = hass.data[DATA_REGISTRY] = Event()
 
         reg = DeviceRegistry(hass)
         await reg.async_load()
 
         hass.data[DATA_REGISTRY] = reg
         evt.set()
-    elif isinstance(reg, Event):
+        return reg
+
+    if isinstance(reg_or_evt, Event):
         # This shouldn't be possible in production (has been tested)
         # It's here soley for testing purposes.
-        await reg.wait()
-        reg = hass.data.get(DATA_REGISTRY)
+        evt = cast(Event, reg_or_evt)
+        await evt.wait()
+        return hass.data.get(DATA_REGISTRY)
 
-    return cast(DeviceRegistry, reg)
+    return cast(DeviceRegistry, reg_or_evt)
 
 
 @callback

diff --git a/homeassistant/helpers/entity_registry.py b/homeassistant/helpers/entity_registry.py
@@ -282,23 +282,26 @@ def async_clear_config_entry(self, config_entry):
 @bind_hass
 async def async_get_registry(hass: HomeAssistantType) -> EntityRegistry:
     """Return entity registry instance."""
-    reg = hass.data.get(DATA_REGISTRY)
+    reg_or_evt = hass.data.get(DATA_REGISTRY)
 
-    if not reg:
+    if not reg_or_evt:
         evt = hass.data[DATA_REGISTRY] = Event()
 
         reg = EntityRegistry(hass)
         await reg.async_load()
 
         hass.data[DATA_REGISTRY] = reg
         evt.set()
-    elif isinstance(reg, Event):
+        return reg
+
+    if isinstance(reg_or_evt, Event):
         # This shouldn't be possible in production (has been tested)
         # It's here soley for testing purposes.
-        await reg.wait()
-        reg = hass.data.get(DATA_REGISTRY)
+        evt = cast(Event, reg_or_evt)
+        await evt.wait()
+        return hass.data.get(DATA_REGISTRY)
 
-    return cast(EntityRegistry, reg)
+    return cast(EntityRegistry, reg_or_evt)
 
 
 @callback