From 32f02aa3c6ac5e56666d5f8941a67398aafb92f0 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Sat, 21 Sep 2024 13:13:41 +0200
Subject: [PATCH] Use HassKey in image (#126322)

---
 homeassistant/components/image/__init__.py     | 10 ++++------
 homeassistant/components/image/const.py        | 13 ++++++++++++-
 homeassistant/components/image/media_source.py | 10 +++-------
 3 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/homeassistant/components/image/__init__.py b/homeassistant/components/image/__init__.py
index 2307a66d5a1..692a398c577 100644
--- a/homeassistant/components/image/__init__.py
+++ b/homeassistant/components/image/__init__.py
@@ -30,7 +30,7 @@ from homeassistant.helpers.event import (
 from homeassistant.helpers.httpx_client import get_async_client
 from homeassistant.helpers.typing import UNDEFINED, ConfigType, UndefinedType
 
-from .const import DOMAIN, IMAGE_TIMEOUT
+from .const import DOMAIN, DOMAIN_DATA, IMAGE_TIMEOUT
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -88,7 +88,7 @@ async def _async_get_image(image_entity: ImageEntity, timeout: int) -> Image:
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up the image component."""
-    component = hass.data[DOMAIN] = EntityComponent[ImageEntity](
+    component = hass.data[DOMAIN_DATA] = EntityComponent[ImageEntity](
         _LOGGER, DOMAIN, hass, SCAN_INTERVAL
     )
 
@@ -120,14 +120,12 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up a config entry."""
-    component: EntityComponent[ImageEntity] = hass.data[DOMAIN]
-    return await component.async_setup_entry(entry)
+    return await hass.data[DOMAIN_DATA].async_setup_entry(entry)
 
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    component: EntityComponent[ImageEntity] = hass.data[DOMAIN]
-    return await component.async_unload_entry(entry)
+    return await hass.data[DOMAIN_DATA].async_unload_entry(entry)
 
 
 CACHED_PROPERTIES_WITH_ATTR_ = {
diff --git a/homeassistant/components/image/const.py b/homeassistant/components/image/const.py
index d96f13b4951..7746e40afbb 100644
--- a/homeassistant/components/image/const.py
+++ b/homeassistant/components/image/const.py
@@ -1,7 +1,18 @@
 """Constants for the image integration."""
 
-from typing import Final
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Final
+
+from homeassistant.util.hass_dict import HassKey
+
+if TYPE_CHECKING:
+    from homeassistant.helpers.entity_component import EntityComponent
+
+    from . import ImageEntity
+
 
 DOMAIN: Final = "image"
+DOMAIN_DATA: HassKey[EntityComponent[ImageEntity]] = HassKey(DOMAIN)
 
 IMAGE_TIMEOUT: Final = 10
diff --git a/homeassistant/components/image/media_source.py b/homeassistant/components/image/media_source.py
index 882249ef940..4ed24498453 100644
--- a/homeassistant/components/image/media_source.py
+++ b/homeassistant/components/image/media_source.py
@@ -14,10 +14,8 @@ from homeassistant.components.media_source import (
 )
 from homeassistant.const import ATTR_FRIENDLY_NAME
 from homeassistant.core import HomeAssistant, State
-from homeassistant.helpers.entity_component import EntityComponent
 
-from . import ImageEntity
-from .const import DOMAIN
+from .const import DOMAIN, DOMAIN_DATA
 
 
 async def async_get_media_source(hass: HomeAssistant) -> ImageMediaSource:
@@ -37,8 +35,7 @@ class ImageMediaSource(MediaSource):
 
     async def async_resolve_media(self, item: MediaSourceItem) -> PlayMedia:
         """Resolve media to a url."""
-        component: EntityComponent[ImageEntity] = self.hass.data[DOMAIN]
-        image = component.get_entity(item.identifier)
+        image = self.hass.data[DOMAIN_DATA].get_entity(item.identifier)
 
         if not image:
             raise Unresolvable(f"Could not resolve media item: {item.identifier}")
@@ -55,7 +52,6 @@ class ImageMediaSource(MediaSource):
         if item.identifier:
             raise BrowseError("Unknown item")
 
-        component: EntityComponent[ImageEntity] = self.hass.data[DOMAIN]
         children = [
             BrowseMediaSource(
                 domain=DOMAIN,
@@ -69,7 +65,7 @@ class ImageMediaSource(MediaSource):
                 can_play=True,
                 can_expand=False,
             )
-            for image in component.entities
+            for image in self.hass.data[DOMAIN_DATA].entities
         ]
 
         return BrowseMediaSource(