Add possibility to configure entoty via labels

custom_components/presence_simulation/__init__.py

@@ -6,6 +6,10 @@
 import pytz
 import re
 import random
+from homeassistant.helpers import (
+        label_registry as lr,
+        entity_registry as er
+)
 from datetime import datetime,timedelta,timezone
 from homeassistant.components.recorder.history import get_significant_states
 from homeassistant.components.recorder import get_instance
@@ -52,6 +56,7 @@ async def stop_presence_simulation(err=None, restart=False, switch_id=None):
             #empty the start_datetime  attribute
             await entity.reset_start_datetime()
             await entity.reset_entities()
+            await entity.reset_labels()
             await entity.reset_delta()
             await entity.reset_random()
             # if the scene exist, turn it on
@@ -89,6 +94,23 @@ async def handle_stop_presence_simulation(call, restart=False, switch_id=None):
         else:
             _LOGGER.warning("Presence simulation switch %s is not on, can't be turned off", switch_id)
 
+    async def async_expand_labels(labels):
+        labels_new = []
+        _LOGGER.debug("expand labels %s", labels)
+        label_reg = lr.async_get(hass)
+        entity_reg = er.async_get(hass)
+        for label_str in labels:
+            _LOGGER.debug("expand label %s", label_str)
+            #to make it asyncable, not sure it is needed
+            await asyncio.sleep(0)
+            if label := label_reg.async_get_label(label_str):
+                _LOGGER.debug("expand label_id %s", label.label_id)
+                for entry in er.async_entries_for_label(entity_reg, label.label_id):
+                    _LOGGER.debug("expand entry %s", entry.entity_id)
+                    labels_new += [entry.entity_id]
+
+        return labels_new
+
     async def async_expand_entities(entities):
         """If the entity is a group, return the list of the entities within, otherwise, return the entity"""
         entities_new = []
@@ -136,6 +158,8 @@ async def handle_presence_simulation(call, restart=False, switch_id=None):
                         await entity.set_entities(call.data.get("entity_id"))
                     else:
                         await entity.set_entities([call.data.get("entity_id")])
+                if "labels" in call.data:
+                    await entity.set_labels(call.data.get("labels"))
                 if "delta" in call.data:
                     await entity.set_delta(call.data.get("delta", 7))
                 if "restore_states" in call.data:
@@ -169,6 +193,12 @@ async def handle_presence_simulation(call, restart=False, switch_id=None):
         except Exception as e:
             _LOGGER.error("Error during identifing entities: "+entity.entities)
             return
+        try:
+            expanded_labels = await async_expand_labels(entity.labels)
+        except Exception as e:
+            _LOGGER.error("Error during identifing labels: "+entity.labels)
+            return
+        expanded_entities+=expanded_labels
 
         if len(expanded_entities) == 0:
             _LOGGER.error("Error during identifing entities, no valid entities has been found")
@@ -514,4 +544,10 @@ def update_unique_id(entity_entry):
         new = {**config_entry.data}
         new["brightness"] = 0
         hass.config_entries.async_update_entry(config_entry, data=new, version=4)
+
+    if config_entry.version == 4: #add labels parameter
+        _LOGGER.debug("Will migrate to version 5")
+        new = {**config_entry.data}
+        new["labels"] = []
+        hass.config_entries.async_update_entry(config_entry, data=new, version=5)
     return True

custom_components/presence_simulation/config_flow.py

@@ -3,6 +3,8 @@
     SelectSelector,
     SelectSelectorConfig,
     SelectSelectorMode,
+    LabelSelector,
+    LabelSelectorConfig
 )
 import re
 import logging
@@ -15,7 +17,7 @@
 _LOGGER = logging.getLogger(__name__)
 
 class PresenceSimulationConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
-    VERSION = 4
+    VERSION = 5
     data = None
     async def async_create_flow(handler, context, data):
             """Create flow."""
@@ -28,6 +30,7 @@ async def async_step_user(self, info=None):
         data_schema = {
             vol.Required("switch", description={"suggested_value": "Choose a unique name"}): str,
             vol.Required("entities"): SelectSelector(SelectSelectorConfig(options=all_entities, multiple=True, mode=SelectSelectorMode.DROPDOWN)),
+            vol.Required("labels"): LabelSelector(LabelSelectorConfig(multiple=True)),
             vol.Required("delta", default=7): int,
             vol.Required("interval", default=30): int,
             vol.Required("restore", default=False): bool,
@@ -51,6 +54,7 @@ async def async_step_user(self, info=None):
         self.data = info
         try:
             _LOGGER.debug("info.entities %s",info['entities'])
+            _LOGGER.debug("info.labels %s",info['labels'])
             #check if entity exist
             #hass.states.get(info['entities'])
         except Exception as e:
@@ -101,6 +105,7 @@ async def async_step_init(self, info=None):
         data_schema = {
             vol.Required("switch", default=self.config_entry.data["switch"]): str,
             vol.Required("entities", default=self.config_entry.data["entities"].split(",")): SelectSelector(SelectSelectorConfig(options=all_entities, multiple=True, mode=SelectSelectorMode.DROPDOWN)),
+            vol.Required("labels", default=self.config_entry.data["labels"]): LabelSelector(LabelSelectorConfig(multiple=True)),
             vol.Required("delta", default=self.config_entry.data["delta"]): int,
             vol.Required("interval", default=interval): int,
             vol.Required("restore", default=restore): bool,

custom_components/presence_simulation/manifest.json

@@ -9,5 +9,5 @@
   "documentation": "https://github.com/slashback100/presence_simulation",
   "iot_class": "local_polling",
   "issue_tracker": "https://github.com/slashback100/presence_simulation/issues",
-  "version": "4.12"
+  "version": "4.13"
 }

custom_components/presence_simulation/switch.py

@@ -60,6 +60,10 @@ def update_config(self, config):
         for elm in conf["entities"].split(","):
             elms += [elm.strip()]
         self._entities = elms
+        if "labels" in conf:
+            self._labels = conf["labels"]
+        else:
+            seld._labels = []
         self._random = int(conf["random"])
         self._interval = int(conf["interval"])
         self._delta = conf["delta"]
@@ -134,6 +138,9 @@ def update(self):
                     del self.attr[prop]
 
     @property
+    def labels(self):
+        return self._labels_overriden
+    @property
     def entities(self):
         return self._entities_overriden
     @property
@@ -157,6 +164,7 @@ def interval(self):
 
     async def reset_default_values_async(self):
         self._entities_overriden = self._entities
+        self._labels_overriden = self._labels
         self._random_overriden = self._random
         self._restore_overriden = self._restore
         self._delta_overriden = self._delta
@@ -165,6 +173,7 @@ async def reset_default_values_async(self):
 
     def reset_default_values(self):
         self._entities_overriden = self._entities
+        self._labels_overriden = self._labels
         self._random_overriden = self._random
         self._restore_overriden = self._restore
         self._delta_overriden = self._delta
@@ -195,6 +204,8 @@ async def async_added_to_hass(self):
                 _LOGGER.debug("State was on")
                 if "entity_id" in state.attributes:
                     self._entities_overriden = state.attributes["entity_id"]
+                if "labels" in state.attributes:
+                    self._labels_overriden = state.attributes["labels"]
                 if "random" in state.attributes:
                     self._random_overriden = state.attributes["random"]
                 if "delta" in state.attributes:
@@ -231,6 +242,11 @@ async def set_delta(self, delta):
         self.attr["delta"] = delta
         self._delta_overriden = delta
 
+    async def set_labels(self, labels):
+        _LOGGER.debug("overidding labels %s", labels)
+        self.attr["labels"] = labels
+        self._labels_overriden = labels
+
     async def set_entities(self, entities):
         _LOGGER.debug("overidding entities %s", entities)
         self.attr["entity_id"] = entities
@@ -264,6 +280,10 @@ async def reset_delta(self):
         if "delta" in self.attr:
             del self.attr["delta"]
 
+    async def reset_labels(self):
+        if "labels" in self.attr:
+            del self.attr["labels"]
+
     async def reset_entities(self):
         if "entity_id" in self.attr:
             del self.attr["entity_id"]