Centralize loading Axis entities (home-assistant#114018)

Centralize platform loading

homeassistant/components/axis/binary_sensor.py

@@ -2,12 +2,11 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable, Iterable
+from collections.abc import Callable
 from dataclasses import dataclass
 from datetime import datetime, timedelta
-from functools import partial
 
-from axis.models.event import Event, EventOperation, EventTopic
+from axis.models.event import Event, EventTopic
 from axis.vapix.interfaces.applications.fence_guard import FenceGuardHandler
 from axis.vapix.interfaces.applications.loitering_guard import LoiteringGuardHandler
 from axis.vapix.interfaces.applications.motion_guard import MotionGuardHandler
@@ -182,28 +181,9 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up a Axis binary sensor."""
-    hub = AxisHub.get_hub(hass, config_entry)
-
-    @callback
-    def register_platform(descriptions: Iterable[AxisBinarySensorDescription]) -> None:
-        """Register entity platform to create entities on event initialized signal."""
-
-        @callback
-        def create_entity(
-            description: AxisBinarySensorDescription, event: Event
-        ) -> None:
-            """Create Axis entity."""
-            if description.supported_fn(hub, event):
-                async_add_entities([AxisBinarySensor(hub, description, event)])
-
-        for description in descriptions:
-            hub.api.event.subscribe(
-                partial(create_entity, description),
-                topic_filter=description.event_topic,
-                operation_filter=EventOperation.INITIALIZED,
-            )
-
-    register_platform(ENTITY_DESCRIPTIONS)
+    AxisHub.get_hub(hass, config_entry).entity_loader.register_platform(
+        async_add_entities, AxisBinarySensor, ENTITY_DESCRIPTIONS
+    )
 
 
 class AxisBinarySensor(AxisEventEntity, BinarySensorEntity):

homeassistant/components/axis/entity.py

@@ -99,8 +99,8 @@ def __init__(
 
         self._event_id = event.id
         self._event_topic = event.topic_base
-        event_type = TOPIC_TO_EVENT_TYPE[event.topic_base]
 
+        event_type = TOPIC_TO_EVENT_TYPE[event.topic_base]
         self._attr_name = description.name_fn(hub, event) or f"{event_type} {event.id}"
 
         self._attr_unique_id = f"{hub.unique_id}-{event.topic}-{event.id}"

homeassistant/components/axis/hub/entity_loader.py

@@ -0,0 +1,72 @@
+"""Axis network device entity loader.
+
+Central point to load entities for the different platforms.
+"""
+from __future__ import annotations
+
+from functools import partial
+from typing import TYPE_CHECKING
+
+from axis.models.event import Event, EventOperation
+
+from homeassistant.core import callback
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+
+from ..entity import AxisEventDescription, AxisEventEntity
+
+if TYPE_CHECKING:
+    from .hub import AxisHub
+
+
+class AxisEntityLoader:
+    """Axis network device integration handling platforms for entity registration."""
+
+    def __init__(self, hub: AxisHub) -> None:
+        """Initialize the UniFi entity loader."""
+        self.hub = hub
+
+        self.platforms: list[
+            tuple[
+                AddEntitiesCallback,
+                type[AxisEventEntity],
+                tuple[AxisEventDescription, ...],
+            ]
+        ] = []
+
+    @callback
+    def register_platform(
+        self,
+        async_add_entities: AddEntitiesCallback,
+        entity_class: type[AxisEventEntity],
+        descriptions: tuple[AxisEventDescription, ...],
+    ) -> None:
+        """Register Axis entity platforms."""
+        self.platforms.append((async_add_entities, entity_class, descriptions))
+
+    @callback
+    def initialize_platforms(self) -> None:
+        """Prepare event listeners and platforms."""
+
+        @callback
+        def load_entities(
+            platform_entity: type[AxisEventEntity],
+            descriptions: tuple[AxisEventDescription, ...],
+            async_add_entities: AddEntitiesCallback,
+        ) -> None:
+            """Set up listeners for events."""
+
+            @callback
+            def create_entity(description: AxisEventDescription, event: Event) -> None:
+                """Create Axis entity."""
+                if description.supported_fn(self.hub, event):
+                    async_add_entities([platform_entity(self.hub, description, event)])
+
+            for description in descriptions:
+                self.hub.api.event.subscribe(
+                    partial(create_entity, description),
+                    topic_filter=description.event_topic,
+                    operation_filter=EventOperation.INITIALIZED,
+                )
+
+        for async_add_entities, entity_class, descriptions in self.platforms:
+            load_entities(entity_class, descriptions, async_add_entities)

homeassistant/components/axis/hub/hub.py

@@ -22,6 +22,7 @@
 
 from ..const import ATTR_MANUFACTURER, DOMAIN as AXIS_DOMAIN
 from .config import AxisConfig
+from .entity_loader import AxisEntityLoader
 
 
 class AxisHub:
@@ -33,6 +34,7 @@ def __init__(
         """Initialize the device."""
         self.hass = hass
         self.config = AxisConfig.from_config_entry(config_entry)
+        self.entity_loader = AxisEntityLoader(self)
         self.api = api
 
         self.available = True
@@ -131,6 +133,8 @@ def mqtt_message(self, message: ReceiveMessage) -> None:
     @callback
     def setup(self) -> None:
         """Set up the device events."""
+        self.entity_loader.initialize_platforms()
+
         self.api.stream.connection_status_callback.append(
             self.connection_status_callback
         )

homeassistant/components/axis/light.py

@@ -1,11 +1,9 @@
 """Support for Axis lights."""
 
-from collections.abc import Iterable
 from dataclasses import dataclass
-from functools import partial
 from typing import Any
 
-from axis.models.event import Event, EventOperation, EventTopic
+from axis.models.event import Event, EventTopic
 
 from homeassistant.components.light import (
     ATTR_BRIGHTNESS,
@@ -51,26 +49,9 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Axis light platform."""
-    hub = AxisHub.get_hub(hass, config_entry)
-
-    @callback
-    def register_platform(descriptions: Iterable[AxisLightDescription]) -> None:
-        """Register entity platform to create entities on event initialized signal."""
-
-        @callback
-        def create_entity(description: AxisLightDescription, event: Event) -> None:
-            """Create Axis entity."""
-            if description.supported_fn(hub, event):
-                async_add_entities([AxisLight(hub, description, event)])
-
-        for description in descriptions:
-            hub.api.event.subscribe(
-                partial(create_entity, description),
-                topic_filter=description.event_topic,
-                operation_filter=EventOperation.INITIALIZED,
-            )
-
-    register_platform(ENTITY_DESCRIPTIONS)
+    AxisHub.get_hub(hass, config_entry).entity_loader.register_platform(
+        async_add_entities, AxisLight, ENTITY_DESCRIPTIONS
+    )
 
 
 class AxisLight(AxisEventEntity, LightEntity):

homeassistant/components/axis/switch.py

@@ -1,11 +1,9 @@
 """Support for Axis switches."""
 
-from collections.abc import Iterable
 from dataclasses import dataclass
-from functools import partial
 from typing import Any
 
-from axis.models.event import Event, EventOperation, EventTopic
+from axis.models.event import Event, EventTopic
 
 from homeassistant.components.switch import (
     SwitchDeviceClass,
@@ -44,26 +42,9 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Axis switch platform."""
-    hub = AxisHub.get_hub(hass, config_entry)
-
-    @callback
-    def register_platform(descriptions: Iterable[AxisSwitchDescription]) -> None:
-        """Register entity platform to create entities on event initialized signal."""
-
-        @callback
-        def create_entity(description: AxisSwitchDescription, event: Event) -> None:
-            """Create Axis entity."""
-            if description.supported_fn(hub, event):
-                async_add_entities([AxisSwitch(hub, description, event)])
-
-        for description in descriptions:
-            hub.api.event.subscribe(
-                partial(create_entity, description),
-                topic_filter=description.event_topic,
-                operation_filter=EventOperation.INITIALIZED,
-            )
-
-    register_platform(ENTITY_DESCRIPTIONS)
+    AxisHub.get_hub(hass, config_entry).entity_loader.register_platform(
+        async_add_entities, AxisSwitch, ENTITY_DESCRIPTIONS
+    )
 
 
 class AxisSwitch(AxisEventEntity, SwitchEntity):