Refactor group setup (#116317)

* Refactor group setup

* Add @callback decorator and remove commented out code

* Keep set, add default on state

homeassistant/components/air_quality/__init__.py

@@ -18,6 +18,7 @@ from homeassistant.helpers.entity_component import EntityComponent
 from homeassistant.helpers.typing import ConfigType, StateType
 
 from . import group as group_pre_import  # noqa: F401
+from .const import DOMAIN
 
 _LOGGER: Final = logging.getLogger(__name__)
 
@@ -33,8 +34,6 @@ ATTR_PM_10: Final = "particulate_matter_10"
 ATTR_PM_2_5: Final = "particulate_matter_2_5"
 ATTR_SO2: Final = "sulphur_dioxide"
 
-DOMAIN: Final = "air_quality"
-
 ENTITY_ID_FORMAT: Final = DOMAIN + ".{}"
 
 SCAN_INTERVAL: Final = timedelta(seconds=30)

homeassistant/components/air_quality/const.py

@@ -0,0 +1,5 @@
+"""Constants for the air_quality entity platform."""
+
+from typing import Final
+
+DOMAIN: Final = "air_quality"

homeassistant/components/air_quality/group.py

@@ -7,10 +7,12 @@ from homeassistant.core import HomeAssistant, callback
 if TYPE_CHECKING:
     from homeassistant.components.group import GroupIntegrationRegistry
 
+from .const import DOMAIN
+
 
 @callback
 def async_describe_on_off_states(
     hass: HomeAssistant, registry: "GroupIntegrationRegistry"
 ) -> None:
     """Describe group on off states."""
-    registry.exclude_domain()
+    registry.exclude_domain(DOMAIN)

homeassistant/components/alarm_control_panel/group.py

@@ -10,9 +10,12 @@ from homeassistant.const import (
     STATE_ALARM_ARMED_VACATION,
     STATE_ALARM_TRIGGERED,
     STATE_OFF,
+    STATE_ON,
 )
 from homeassistant.core import HomeAssistant, callback
 
+from .const import DOMAIN
+
 if TYPE_CHECKING:
     from homeassistant.components.group import GroupIntegrationRegistry
 
@@ -23,7 +26,9 @@ def async_describe_on_off_states(
 ) -> None:
     """Describe group on off states."""
     registry.on_off_states(
+        DOMAIN,
         {
+            STATE_ON,
             STATE_ALARM_ARMED_AWAY,
             STATE_ALARM_ARMED_CUSTOM_BYPASS,
             STATE_ALARM_ARMED_HOME,
@@ -31,5 +36,6 @@ def async_describe_on_off_states(
             STATE_ALARM_ARMED_VACATION,
             STATE_ALARM_TRIGGERED,
         },
+        STATE_ON,
         STATE_OFF,
     )

homeassistant/components/climate/group.py

@@ -2,10 +2,10 @@
 
 from typing import TYPE_CHECKING
 
-from homeassistant.const import STATE_OFF
+from homeassistant.const import STATE_OFF, STATE_ON
 from homeassistant.core import HomeAssistant, callback
 
-from .const import HVAC_MODES, HVACMode
+from .const import DOMAIN, HVACMode
 
 if TYPE_CHECKING:
     from homeassistant.components.group import GroupIntegrationRegistry
@@ -17,6 +17,15 @@ def async_describe_on_off_states(
 ) -> None:
     """Describe group on off states."""
     registry.on_off_states(
-        set(HVAC_MODES) - {HVACMode.OFF},
+        DOMAIN,
+        {
+            STATE_ON,
+            HVACMode.HEAT,
+            HVACMode.COOL,
+            HVACMode.HEAT_COOL,
+            HVACMode.AUTO,
+            HVACMode.FAN_ONLY,
+        },
+        STATE_ON,
         STATE_OFF,
     )

homeassistant/components/cover/__init__.py

@@ -46,10 +46,10 @@ from homeassistant.helpers.typing import ConfigType
 from homeassistant.loader import bind_hass
 
 from . import group as group_pre_import  # noqa: F401
+from .const import DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
-DOMAIN = "cover"
 SCAN_INTERVAL = timedelta(seconds=15)
 
 ENTITY_ID_FORMAT = DOMAIN + ".{}"

homeassistant/components/cover/const.py

@@ -0,0 +1,3 @@
+"""Constants for cover entity platform."""
+
+DOMAIN = "cover"

homeassistant/components/cover/group.py

@@ -5,6 +5,8 @@ from typing import TYPE_CHECKING
 from homeassistant.const import STATE_CLOSED, STATE_OPEN
 from homeassistant.core import HomeAssistant, callback
 
+from .const import DOMAIN
+
 if TYPE_CHECKING:
     from homeassistant.components.group import GroupIntegrationRegistry
 
@@ -15,4 +17,4 @@ def async_describe_on_off_states(
 ) -> None:
     """Describe group on off states."""
     # On means open, Off means closed
-    registry.on_off_states({STATE_OPEN}, STATE_CLOSED)
+    registry.on_off_states(DOMAIN, {STATE_OPEN}, STATE_OPEN, STATE_CLOSED)

homeassistant/components/device_tracker/group.py

@@ -5,6 +5,8 @@ from typing import TYPE_CHECKING
 from homeassistant.const import STATE_HOME, STATE_NOT_HOME
 from homeassistant.core import HomeAssistant, callback
 
+from .const import DOMAIN
+
 if TYPE_CHECKING:
     from homeassistant.components.group import GroupIntegrationRegistry
 
@@ -14,4 +16,4 @@ def async_describe_on_off_states(
     hass: HomeAssistant, registry: "GroupIntegrationRegistry"
 ) -> None:
     """Describe group on off states."""
-    registry.on_off_states({STATE_HOME}, STATE_NOT_HOME)
+    registry.on_off_states(DOMAIN, {STATE_HOME}, STATE_HOME, STATE_NOT_HOME)

homeassistant/components/group/registry.py

@@ -2,8 +2,7 @@
 
 from __future__ import annotations
 
-from contextvars import ContextVar
-from typing import Protocol
+from typing import TYPE_CHECKING, Protocol
 
 from homeassistant.const import STATE_OFF, STATE_ON
 from homeassistant.core import HomeAssistant, callback
@@ -13,12 +12,13 @@ from homeassistant.helpers.integration_platform import (
 
 from .const import DOMAIN, REG_KEY
 
-current_domain: ContextVar[str] = ContextVar("current_domain")
+if TYPE_CHECKING:
+    from .entity import Group
 
 
 async def async_setup(hass: HomeAssistant) -> None:
     """Set up the Group integration registry of integration platforms."""
-    hass.data[REG_KEY] = GroupIntegrationRegistry()
+    hass.data[REG_KEY] = GroupIntegrationRegistry(hass)
 
     await async_process_integration_platforms(
         hass, DOMAIN, _process_group_platform, wait_for_platforms=True
@@ -39,7 +39,6 @@ def _process_group_platform(
     hass: HomeAssistant, domain: str, platform: GroupProtocol
 ) -> None:
     """Process a group platform."""
-    current_domain.set(domain)
     registry: GroupIntegrationRegistry = hass.data[REG_KEY]
     platform.async_describe_on_off_states(hass, registry)
 
@@ -47,24 +46,31 @@ def _process_group_platform(
 class GroupIntegrationRegistry:
     """Class to hold a registry of integrations."""
 
-    def __init__(self) -> None:
+    def __init__(self, hass: HomeAssistant) -> None:
         """Imitialize registry."""
+        self.hass = hass
         self.on_off_mapping: dict[str, str] = {STATE_ON: STATE_OFF}
         self.off_on_mapping: dict[str, str] = {STATE_OFF: STATE_ON}
         self.on_states_by_domain: dict[str, set[str]] = {}
         self.exclude_domains: set[str] = set()
+        self.state_group_mapping: dict[str, tuple[str, str]] = {}
+        self.group_entities: set[Group] = set()
 
-    def exclude_domain(self) -> None:
+    @callback
+    def exclude_domain(self, domain: str) -> None:
         """Exclude the current domain."""
-        self.exclude_domains.add(current_domain.get())
+        self.exclude_domains.add(domain)
 
-    def on_off_states(self, on_states: set, off_state: str) -> None:
+    @callback
+    def on_off_states(
+        self, domain: str, on_states: set[str], default_on_state: str, off_state: str
+    ) -> None:
         """Register on and off states for the current domain."""
         for on_state in on_states:
             if on_state not in self.on_off_mapping:
                 self.on_off_mapping[on_state] = off_state
 
         if len(on_states) == 1 and off_state not in self.off_on_mapping:
-            self.off_on_mapping[off_state] = list(on_states)[0]
+            self.off_on_mapping[off_state] = default_on_state
 
-        self.on_states_by_domain[current_domain.get()] = set(on_states)
+        self.on_states_by_domain[domain] = on_states

homeassistant/components/lock/__init__.py

@@ -44,13 +44,13 @@ from homeassistant.helpers.entity_component import EntityComponent
 from homeassistant.helpers.typing import ConfigType, StateType
 
 from . import group as group_pre_import  # noqa: F401
+from .const import DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
 ATTR_CHANGED_BY = "changed_by"
 CONF_DEFAULT_CODE = "default_code"
 
-DOMAIN = "lock"
 SCAN_INTERVAL = timedelta(seconds=30)
 
 ENTITY_ID_FORMAT = DOMAIN + ".{}"

homeassistant/components/lock/const.py

@@ -0,0 +1,3 @@
+"""Constants for the lock entity platform."""
+
+DOMAIN = "lock"

homeassistant/components/lock/group.py

@@ -5,6 +5,8 @@ from typing import TYPE_CHECKING
 from homeassistant.const import STATE_LOCKED, STATE_UNLOCKED
 from homeassistant.core import HomeAssistant, callback
 
+from .const import DOMAIN
+
 if TYPE_CHECKING:
     from homeassistant.components.group import GroupIntegrationRegistry
 
@@ -14,4 +16,4 @@ def async_describe_on_off_states(
     hass: HomeAssistant, registry: "GroupIntegrationRegistry"
 ) -> None:
     """Describe group on off states."""
-    registry.on_off_states({STATE_UNLOCKED}, STATE_LOCKED)
+    registry.on_off_states(DOMAIN, {STATE_UNLOCKED}, STATE_UNLOCKED, STATE_LOCKED)

homeassistant/components/media_player/group.py

@@ -11,6 +11,8 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant, callback
 
+from .const import DOMAIN
+
 if TYPE_CHECKING:
     from homeassistant.components.group import GroupIntegrationRegistry
 
@@ -21,5 +23,13 @@ def async_describe_on_off_states(
 ) -> None:
     """Describe group on off states."""
     registry.on_off_states(
-        {STATE_ON, STATE_PAUSED, STATE_PLAYING, STATE_IDLE}, STATE_OFF
+        DOMAIN,
+        {
+            STATE_ON,
+            STATE_PAUSED,
+            STATE_PLAYING,
+            STATE_IDLE,
+        },
+        STATE_ON,
+        STATE_OFF,
     )

homeassistant/components/person/__init__.py

@@ -55,6 +55,7 @@ from homeassistant.helpers.typing import ConfigType
 from homeassistant.loader import bind_hass
 
 from . import group as group_pre_import  # noqa: F401
+from .const import DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -66,8 +67,6 @@ CONF_DEVICE_TRACKERS = "device_trackers"
 CONF_USER_ID = "user_id"
 CONF_PICTURE = "picture"
 
-DOMAIN = "person"
-
 STORAGE_KEY = DOMAIN
 STORAGE_VERSION = 2
 # Device tracker states to ignore

homeassistant/components/person/const.py

@@ -0,0 +1,3 @@
+"""Constants for the person entity platform."""
+
+DOMAIN = "person"

homeassistant/components/person/group.py

@@ -5,6 +5,8 @@ from typing import TYPE_CHECKING
 from homeassistant.const import STATE_HOME, STATE_NOT_HOME
 from homeassistant.core import HomeAssistant, callback
 
+from .const import DOMAIN
+
 if TYPE_CHECKING:
     from homeassistant.components.group import GroupIntegrationRegistry
 
@@ -14,4 +16,4 @@ def async_describe_on_off_states(
     hass: HomeAssistant, registry: "GroupIntegrationRegistry"
 ) -> None:
     """Describe group on off states."""
-    registry.on_off_states({STATE_HOME}, STATE_NOT_HOME)
+    registry.on_off_states(DOMAIN, {STATE_HOME}, STATE_HOME, STATE_NOT_HOME)

homeassistant/components/plant/group.py

@@ -5,6 +5,8 @@ from typing import TYPE_CHECKING
 from homeassistant.const import STATE_OK, STATE_PROBLEM
 from homeassistant.core import HomeAssistant, callback
 
+from .const import DOMAIN
+
 if TYPE_CHECKING:
     from homeassistant.components.group import GroupIntegrationRegistry
 
@@ -14,4 +16,4 @@ def async_describe_on_off_states(
     hass: HomeAssistant, registry: "GroupIntegrationRegistry"
 ) -> None:
     """Describe group on off states."""
-    registry.on_off_states({STATE_PROBLEM}, STATE_OK)
+    registry.on_off_states(DOMAIN, {STATE_PROBLEM}, STATE_PROBLEM, STATE_OK)

homeassistant/components/sensor/group.py

@@ -7,10 +7,12 @@ from homeassistant.core import HomeAssistant, callback
 if TYPE_CHECKING:
     from homeassistant.components.group import GroupIntegrationRegistry
 
+from .const import DOMAIN
+
 
 @callback
 def async_describe_on_off_states(
     hass: HomeAssistant, registry: "GroupIntegrationRegistry"
 ) -> None:
     """Describe group on off states."""
-    registry.exclude_domain()
+    registry.exclude_domain(DOMAIN)

homeassistant/components/vacuum/__init__.py

@@ -36,11 +36,10 @@ from homeassistant.helpers.typing import ConfigType
 from homeassistant.loader import bind_hass
 
 from . import group as group_pre_import  # noqa: F401
-from .const import STATE_CLEANING, STATE_DOCKED, STATE_ERROR, STATE_RETURNING
+from .const import DOMAIN, STATE_CLEANING, STATE_DOCKED, STATE_ERROR, STATE_RETURNING
 
 _LOGGER = logging.getLogger(__name__)
 
-DOMAIN = "vacuum"
 ENTITY_ID_FORMAT = DOMAIN + ".{}"
 SCAN_INTERVAL = timedelta(seconds=20)
 

homeassistant/components/vacuum/const.py

@@ -1,5 +1,7 @@
 """Support for vacuum cleaner robots (botvacs)."""
 
+DOMAIN = "vacuum"
+
 STATE_CLEANING = "cleaning"
 STATE_DOCKED = "docked"
 STATE_RETURNING = "returning"

homeassistant/components/vacuum/group.py

@@ -7,7 +7,8 @@ from homeassistant.core import HomeAssistant, callback
 
 if TYPE_CHECKING:
     from homeassistant.components.group import GroupIntegrationRegistry
-from .const import STATE_CLEANING, STATE_ERROR, STATE_RETURNING
+
+from .const import DOMAIN, STATE_CLEANING, STATE_ERROR, STATE_RETURNING
 
 
 @callback
@@ -16,5 +17,13 @@ def async_describe_on_off_states(
 ) -> None:
     """Describe group on off states."""
     registry.on_off_states(
-        {STATE_CLEANING, STATE_ON, STATE_RETURNING, STATE_ERROR}, STATE_OFF
+        DOMAIN,
+        {
+            STATE_ON,
+            STATE_CLEANING,
+            STATE_RETURNING,
+            STATE_ERROR,
+        },
+        STATE_ON,
+        STATE_OFF,
     )

homeassistant/components/water_heater/__init__.py

@@ -44,12 +44,11 @@ from homeassistant.helpers.typing import ConfigType
 from homeassistant.util.unit_conversion import TemperatureConverter
 
 from . import group as group_pre_import  # noqa: F401
+from .const import DOMAIN
 
 DEFAULT_MIN_TEMP = 110
 DEFAULT_MAX_TEMP = 140
 
-DOMAIN = "water_heater"
-
 ENTITY_ID_FORMAT = DOMAIN + ".{}"
 SCAN_INTERVAL = timedelta(seconds=60)
 

homeassistant/components/water_heater/const.py

@@ -1,5 +1,7 @@
 """Support for water heater devices."""
 
+DOMAIN = "water_heater"
+
 STATE_ECO = "eco"
 STATE_ELECTRIC = "electric"
 STATE_PERFORMANCE = "performance"

homeassistant/components/water_heater/group.py

@@ -2,12 +2,14 @@
 
 from typing import TYPE_CHECKING
 
-from homeassistant.const import STATE_OFF
+from homeassistant.const import STATE_OFF, STATE_ON
 from homeassistant.core import HomeAssistant, callback
 
 if TYPE_CHECKING:
     from homeassistant.components.group import GroupIntegrationRegistry
+
 from .const import (
+    DOMAIN,
     STATE_ECO,
     STATE_ELECTRIC,
     STATE_GAS,
@@ -23,7 +25,9 @@ def async_describe_on_off_states(
 ) -> None:
     """Describe group on off states."""
     registry.on_off_states(
+        DOMAIN,
         {
+            STATE_ON,
             STATE_ECO,
             STATE_ELECTRIC,
             STATE_PERFORMANCE,
@@ -31,5 +35,6 @@ def async_describe_on_off_states(
             STATE_HEAT_PUMP,
             STATE_GAS,
         },
+        STATE_ON,
         STATE_OFF,
     )

homeassistant/components/weather/group.py

@@ -7,10 +7,12 @@ from homeassistant.core import HomeAssistant, callback
 if TYPE_CHECKING:
     from homeassistant.components.group import GroupIntegrationRegistry
 
+from .const import DOMAIN
+
 
 @callback
 def async_describe_on_off_states(
     hass: HomeAssistant, registry: "GroupIntegrationRegistry"
 ) -> None:
     """Describe group on off states."""
-    registry.exclude_domain()
+    registry.exclude_domain(DOMAIN)