Use HassKey in group (#126321)

* Use HassKey in group

* Adjust

* Improve

homeassistant/components/group/__init__.py

@@ -22,7 +22,6 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant, ServiceCall
 from homeassistant.helpers import config_validation as cv, entity_registry as er
-from homeassistant.helpers.entity_component import EntityComponent
 from homeassistant.helpers.group import (
     expand_entity_ids as _expand_entity_ids,
     get_entity_ids as _get_entity_ids,
@@ -50,11 +49,12 @@ from .const import (  # noqa: F401
     ATTR_REMOVE_ENTITIES,
     CONF_HIDE_MEMBERS,
     DOMAIN,
+    DOMAIN_DATA,
     GROUP_ORDER,
     REG_KEY,
 )
 from .entity import Group, async_get_component
-from .registry import GroupIntegrationRegistry, async_setup as async_setup_registry
+from .registry import async_setup as async_setup_registry
 
 CONF_ALL = "all"
 
@@ -110,8 +110,7 @@ def is_on(hass: HomeAssistant, entity_id: str) -> bool:
         return False
 
     if (state := hass.states.get(entity_id)) is not None:
-        registry: GroupIntegrationRegistry = hass.data[REG_KEY]
-        return state.state in registry.on_off_mapping
+        return state.state in hass.data[REG_KEY].on_off_mapping
 
     return False
 
@@ -132,7 +131,7 @@ def groups_with_entity(hass: HomeAssistant, entity_id: str) -> list[str]:
 
     return [
         group.entity_id
-        for group in hass.data[DOMAIN].entities
+        for group in hass.data[DOMAIN_DATA].entities
         if entity_id in group.tracking
     ]
 
@@ -179,10 +178,7 @@ async def async_remove_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up all groups found defined in the configuration."""
-    if DOMAIN not in hass.data:
-        hass.data[DOMAIN] = EntityComponent[Group](_LOGGER, DOMAIN, hass)
-
-    component: EntityComponent[Group] = hass.data[DOMAIN]
+    component = async_get_component(hass)
 
     await async_setup_registry(hass)
 
@@ -338,7 +334,7 @@ async def _async_process_config(hass: HomeAssistant, config: ConfigType) -> None
         entity_ids: Collection[str] = conf.get(CONF_ENTITIES) or []
         icon: str | None = conf.get(CONF_ICON)
         mode = bool(conf.get(CONF_ALL))
-        order: int = hass.data[GROUP_ORDER]
+        order = hass.data[GROUP_ORDER]
 
         # We keep track of the order when we are creating the tasks
         # in the same way that async_create_group does to make

homeassistant/components/group/const.py

@@ -1,14 +1,24 @@
 """Constants for the Group integration."""
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from homeassistant.util.hass_dict import HassKey
+
+if TYPE_CHECKING:
+    from homeassistant.helpers.entity_component import EntityComponent
+
+    from .entity import Group
+    from .registry import GroupIntegrationRegistry
+
 CONF_HIDE_MEMBERS = "hide_members"
 CONF_IGNORE_NON_NUMERIC = "ignore_non_numeric"
 
 DOMAIN = "group"
-
-REG_KEY = f"{DOMAIN}_registry"
-
-GROUP_ORDER = "group_order"
-
+DOMAIN_DATA: HassKey[EntityComponent[Group]] = HassKey(DOMAIN)
+REG_KEY: HassKey[GroupIntegrationRegistry] = HassKey(f"{DOMAIN}_registry")
+GROUP_ORDER: HassKey[int] = HassKey("group_order")
 
 ATTR_ADD_ENTITIES = "add_entities"
 ATTR_REMOVE_ENTITIES = "remove_entities"

homeassistant/components/group/entity.py

@@ -22,7 +22,7 @@ from homeassistant.helpers.entity import Entity, async_generate_entity_id
 from homeassistant.helpers.entity_component import EntityComponent
 from homeassistant.helpers.event import async_track_state_change_event
 
-from .const import ATTR_AUTO, ATTR_ORDER, DOMAIN, GROUP_ORDER, REG_KEY
+from .const import ATTR_AUTO, ATTR_ORDER, DOMAIN, DOMAIN_DATA, GROUP_ORDER, REG_KEY
 from .registry import GroupIntegrationRegistry, SingleStateType
 
 ENTITY_ID_FORMAT = DOMAIN + ".{}"
@@ -478,8 +478,8 @@ class Group(Entity):
 
 def async_get_component(hass: HomeAssistant) -> EntityComponent[Group]:
     """Get the group entity component."""
-    if (component := hass.data.get(DOMAIN)) is None:
-        component = hass.data[DOMAIN] = EntityComponent[Group](
+    if (component := hass.data.get(DOMAIN_DATA)) is None:
+        component = hass.data[DOMAIN_DATA] = EntityComponent[Group](
             _PACKAGE_LOGGER, DOMAIN, hass
         )
     return component

homeassistant/components/group/registry.py

@@ -160,8 +160,7 @@ def _process_group_platform(
     hass: HomeAssistant, domain: str, platform: GroupProtocol
 ) -> None:
     """Process a group platform."""
-    registry: GroupIntegrationRegistry = hass.data[REG_KEY]
-    platform.async_describe_on_off_states(hass, registry)
+    platform.async_describe_on_off_states(hass, hass.data[REG_KEY])
 
 
 @dataclass(frozen=True, slots=True)