diff --git a/homeassistant/components/ai_task/__init__.py b/homeassistant/components/ai_task/__init__.py index e88a2960379f..3840fc515b0a 100644 --- a/homeassistant/components/ai_task/__init__.py +++ b/homeassistant/components/ai_task/__init__.py @@ -5,6 +5,7 @@ from typing import Any import voluptuous as vol +from homeassistant.components.media_source import local_source from homeassistant.config_entries import ConfigEntry from homeassistant.const import ATTR_ENTITY_ID, CONF_DESCRIPTION, CONF_SELECTOR from homeassistant.core import ( @@ -34,6 +35,7 @@ from .const import ( ) from .entity import AITaskEntity from .http import async_setup as async_setup_http +from .media_source import async_get_media_source from .task import ( GenDataTask, GenDataTaskResult, @@ -88,6 +90,9 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: hass.data[DATA_PREFERENCES] = AITaskPreferences(hass) await hass.data[DATA_PREFERENCES].async_load() async_setup_http(hass) + if hass.config.media_dirs: + source = await async_get_media_source(hass) + hass.http.register_view(local_source.LocalMediaView(hass, source)) hass.services.async_register( DOMAIN, SERVICE_GENERATE_DATA, diff --git a/homeassistant/components/ai_task/media_source.py b/homeassistant/components/ai_task/media_source.py index 9f0e493b0ada..378809eee554 100644 --- a/homeassistant/components/ai_task/media_source.py +++ b/homeassistant/components/ai_task/media_source.py @@ -2,14 +2,16 @@ from pathlib import Path -from homeassistant.components.media_source import MediaSource, local_source +from homeassistant.components.media_source import local_source from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError +from homeassistant.helpers.singleton import singleton from .const import DATA_MEDIA_SOURCE, DOMAIN, IMAGE_DIR -async def async_get_media_source(hass: HomeAssistant) -> MediaSource: +@singleton(DATA_MEDIA_SOURCE, async_=True) +async def async_get_media_source(hass: HomeAssistant) -> local_source.LocalSource: """Set up local media source.""" media_dirs = list(hass.config.media_dirs.values()) @@ -20,11 +22,10 @@ async def async_get_media_source(hass: HomeAssistant) -> MediaSource: media_dir = Path(media_dirs[0]) / DOMAIN / IMAGE_DIR - hass.data[DATA_MEDIA_SOURCE] = source = local_source.LocalSource( + return local_source.LocalSource( hass, DOMAIN, "AI generated images", {IMAGE_DIR: str(media_dir)}, f"/{DOMAIN}", ) - return source diff --git a/homeassistant/components/media_source/__init__.py b/homeassistant/components/media_source/__init__.py index 030f68bf4142..bc68335fcad3 100644 --- a/homeassistant/components/media_source/__init__.py +++ b/homeassistant/components/media_source/__init__.py @@ -5,17 +5,16 @@ from typing import Protocol from homeassistant.components import websocket_api from homeassistant.core import HomeAssistant from homeassistant.helpers import config_validation as cv -from homeassistant.helpers.integration_platform import ( - async_process_integration_platforms, -) +from homeassistant.helpers.integration_platform import LazyIntegrationPlatforms from homeassistant.helpers.typing import ConfigType from . import http, local_source from .const import ( + DATA_LOCAL_SOURCE, + DATA_MEDIA_SOURCE_PLATFORMS, DOMAIN, MEDIA_CLASS_MAP, MEDIA_MIME_TYPES, - MEDIA_SOURCE_DATA, URI_SCHEME, URI_SCHEME_REGEX, ) @@ -73,17 +72,18 @@ def generate_media_source_id(domain: str, identifier: str) -> str: async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: """Set up the media_source component.""" - hass.data[MEDIA_SOURCE_DATA] = {} + hass.data[DATA_MEDIA_SOURCE_PLATFORMS] = LazyIntegrationPlatforms[MediaSource]( + hass, DOMAIN, _process_media_source_platform + ) http.async_setup(hass) # Local sources support - await _process_media_source_platform(hass, DOMAIN, local_source) + source = await local_source.async_get_media_source(hass) + hass.data[DATA_LOCAL_SOURCE] = source + hass.http.register_view(local_source.LocalMediaView(hass, source)) hass.http.register_view(local_source.UploadMediaView) websocket_api.async_register_command(hass, local_source.websocket_remove_media) - await async_process_integration_platforms( - hass, DOMAIN, _process_media_source_platform - ) return True @@ -91,9 +91,6 @@ async def _process_media_source_platform( hass: HomeAssistant, domain: str, platform: MediaSourceProtocol, -) -> None: +) -> MediaSource: """Process a media source platform.""" - source = await platform.async_get_media_source(hass) - hass.data[MEDIA_SOURCE_DATA][domain] = source - if isinstance(source, local_source.LocalSource): - hass.http.register_view(local_source.LocalMediaView(hass, source)) + return await platform.async_get_media_source(hass) diff --git a/homeassistant/components/media_source/const.py b/homeassistant/components/media_source/const.py index 1e9a7cc1eaa2..e7cc4499b006 100644 --- a/homeassistant/components/media_source/const.py +++ b/homeassistant/components/media_source/const.py @@ -7,10 +7,15 @@ from homeassistant.components.media_player import MediaClass from homeassistant.util.hass_dict import HassKey if TYPE_CHECKING: + from homeassistant.helpers.integration_platform import LazyIntegrationPlatforms + from .models import MediaSource DOMAIN = "media_source" -MEDIA_SOURCE_DATA: HassKey[dict[str, MediaSource]] = HassKey(DOMAIN) +DATA_LOCAL_SOURCE: HassKey[MediaSource] = HassKey("media_source_local_source") +DATA_MEDIA_SOURCE_PLATFORMS: HassKey[LazyIntegrationPlatforms[MediaSource]] = HassKey( + "media_source_platforms" +) MEDIA_MIME_TYPES = ("audio", "video", "image") MEDIA_CLASS_MAP = { "audio": MediaClass.MUSIC, diff --git a/homeassistant/components/media_source/helper.py b/homeassistant/components/media_source/helper.py index 099774deaa8d..92eff623bc3d 100644 --- a/homeassistant/components/media_source/helper.py +++ b/homeassistant/components/media_source/helper.py @@ -8,17 +8,23 @@ from homeassistant.components.media_player import ( SearchMedia, SearchMediaQuery, ) -from homeassistant.core import HomeAssistant, callback +from homeassistant.core import HomeAssistant from homeassistant.helpers.frame import report_usage from homeassistant.helpers.typing import UNDEFINED, UndefinedType -from .const import DOMAIN, MEDIA_SOURCE_DATA +from .const import DOMAIN from .error import UnknownMediaSource, Unresolvable -from .models import BrowseMediaSource, MediaSourceItem, PlayMedia, RootBrowseMediaSource +from .models import ( + BrowseMediaSource, + MediaSourceItem, + PlayMedia, + RootBrowseMediaSource, + _async_get_media_source, + _async_get_media_sources, +) -@callback -def _get_media_item( +async def _get_media_item( hass: HomeAssistant, media_content_id: str | None, target_media_player: str | None ) -> MediaSourceItem: """Return media item.""" @@ -26,10 +32,14 @@ def _get_media_item( item = MediaSourceItem.from_uri(hass, media_content_id, target_media_player) else: # We default to our own domain if its only one registered - domain = None if len(hass.data[MEDIA_SOURCE_DATA]) > 1 else DOMAIN + sources = await _async_get_media_sources(hass) + domain = None if len(sources) > 1 else DOMAIN return MediaSourceItem(hass, domain, "", target_media_player) - if item.domain is not None and item.domain not in hass.data[MEDIA_SOURCE_DATA]: + if ( + item.domain is not None + and await _async_get_media_source(hass, item.domain) is None + ): raise UnknownMediaSource( translation_domain=DOMAIN, translation_key="unknown_media_source", @@ -46,11 +56,12 @@ async def async_browse_media( content_filter: Callable[[BrowseMedia], bool] | None = None, ) -> BrowseMediaSource | RootBrowseMediaSource: """Return media player browse media results.""" - if DOMAIN not in hass.data: + if DOMAIN not in hass.config.top_level_components: raise BrowseError("Media Source not loaded") try: - item = await _get_media_item(hass, media_content_id, None).async_browse() + media_item = await _get_media_item(hass, media_content_id, None) + item = await media_item.async_browse() except ValueError as err: raise BrowseError( translation_domain=DOMAIN, @@ -78,11 +89,12 @@ async def async_search_media( query: SearchMediaQuery, ) -> SearchMedia: """Return media searched in the media source.""" - if DOMAIN not in hass.data: + if DOMAIN not in hass.config.top_level_components: raise BrowseError("Media Source not loaded") try: - return await _get_media_item(hass, media_content_id, None).async_search(query) + media_item = await _get_media_item(hass, media_content_id, None) + return await media_item.async_search(query) except NotImplementedError as err: raise BrowseError( translation_domain=DOMAIN, @@ -106,7 +118,7 @@ async def async_resolve_media( target_media_player: str | None | UndefinedType = UNDEFINED, ) -> PlayMedia: """Get info to play media.""" - if DOMAIN not in hass.data: + if DOMAIN not in hass.config.top_level_components: raise Unresolvable("Media Source not loaded") if target_media_player is UNDEFINED: @@ -117,7 +129,7 @@ async def async_resolve_media( target_media_player = None try: - item = _get_media_item(hass, media_content_id, target_media_player) + item = await _get_media_item(hass, media_content_id, target_media_player) except ValueError as err: raise Unresolvable( translation_domain=DOMAIN, diff --git a/homeassistant/components/media_source/local_source.py b/homeassistant/components/media_source/local_source.py index 5d9ca3fb581e..49fb54b225ae 100644 --- a/homeassistant/components/media_source/local_source.py +++ b/homeassistant/components/media_source/local_source.py @@ -24,7 +24,7 @@ from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import HomeAssistantError from homeassistant.util import raise_if_invalid_filename, raise_if_invalid_path -from .const import DOMAIN, MEDIA_CLASS_MAP, MEDIA_MIME_TYPES, MEDIA_SOURCE_DATA +from .const import DATA_LOCAL_SOURCE, DOMAIN, MEDIA_CLASS_MAP, MEDIA_MIME_TYPES from .error import Unresolvable from .models import BrowseMediaSource, MediaSource, MediaSourceItem, PlayMedia @@ -446,7 +446,7 @@ class UploadMediaView(http.HomeAssistantView): if target_folder.domain != DOMAIN: raise web.HTTPBadRequest - source = cast(LocalSource, hass.data[MEDIA_SOURCE_DATA][target_folder.domain]) + source = cast(LocalSource, hass.data[DATA_LOCAL_SOURCE]) try: uploaded_media_source_id = await source.async_upload_media( target_folder, data["file"] @@ -491,7 +491,7 @@ async def websocket_remove_media( ) return - source = cast(LocalSource, hass.data[MEDIA_SOURCE_DATA][item.domain]) + source = cast(LocalSource, hass.data[DATA_LOCAL_SOURCE]) try: await source.async_delete_media(item) diff --git a/homeassistant/components/media_source/models.py b/homeassistant/components/media_source/models.py index fdf34c9f594b..2a341f3d0c3b 100644 --- a/homeassistant/components/media_source/models.py +++ b/homeassistant/components/media_source/models.py @@ -10,15 +10,37 @@ from homeassistant.components.media_player import ( SearchMedia, SearchMediaQuery, ) -from homeassistant.core import HomeAssistant, callback +from homeassistant.core import HomeAssistant from homeassistant.helpers.translation import async_get_cached_translations -from .const import MEDIA_SOURCE_DATA, URI_SCHEME, URI_SCHEME_REGEX +from .const import ( + DATA_LOCAL_SOURCE, + DATA_MEDIA_SOURCE_PLATFORMS, + DOMAIN, + URI_SCHEME, + URI_SCHEME_REGEX, +) if TYPE_CHECKING: from pathlib import Path +async def _async_get_media_sources(hass: HomeAssistant) -> dict[str, MediaSource]: + """Return all media sources, loading integration platforms on demand.""" + sources: dict[str, MediaSource] = {DOMAIN: hass.data[DATA_LOCAL_SOURCE]} + sources.update(await hass.data[DATA_MEDIA_SOURCE_PLATFORMS].async_get_platforms()) + return sources + + +async def _async_get_media_source( + hass: HomeAssistant, domain: str +) -> MediaSource | None: + """Return the media source for a domain, loading it on demand.""" + if domain == DOMAIN: + return hass.data[DATA_LOCAL_SOURCE] + return await hass.data[DATA_MEDIA_SOURCE_PLATFORMS].async_get_platform(domain) + + @dataclass(slots=True) class PlayMedia: """Represents a playable media.""" @@ -87,6 +109,7 @@ class MediaSourceItem: can_expand=True, children_media_class=MediaClass.APP, ) + sources = await _async_get_media_sources(self.hass) base.children = sorted( ( BrowseMediaSource( @@ -99,13 +122,14 @@ class MediaSourceItem: can_play=False, can_expand=True, ) - for source in self.hass.data[MEDIA_SOURCE_DATA].values() + for source in sources.values() ), key=lambda item: item.title, ) return base - return await self.async_media_source().async_browse_media(self) + source = await self._async_media_source() + return await source.async_browse_media(self) async def async_search(self, query: SearchMediaQuery) -> SearchMedia: """Search this item.""" @@ -114,18 +138,21 @@ class MediaSourceItem: if self.domain is None: raise NotImplementedError - return await self.async_media_source().async_search_media(self, query) + return await (await self._async_media_source()).async_search_media(self, query) async def async_resolve(self) -> PlayMedia: """Resolve to playable item.""" - return await self.async_media_source().async_resolve_media(self) + source = await self._async_media_source() + return await source.async_resolve_media(self) - @callback - def async_media_source(self) -> MediaSource: + async def _async_media_source(self) -> MediaSource: """Return media source that owns this item.""" if TYPE_CHECKING: assert self.domain is not None - return self.hass.data[MEDIA_SOURCE_DATA][self.domain] + # Existence is validated by _get_media_item before browse/resolve. + source = await _async_get_media_source(self.hass, self.domain) + assert source is not None + return source @classmethod def from_uri( diff --git a/homeassistant/helpers/integration_platform.py b/homeassistant/helpers/integration_platform.py index 7fc212c9dba3..236eb7556138 100644 --- a/homeassistant/helpers/integration_platform.py +++ b/homeassistant/helpers/integration_platform.py @@ -266,7 +266,7 @@ async def _async_process_integration_platforms( # Any = platform. -type ProcessPlatform[_R] = Callable[[HomeAssistant, str, Any], _R] +type ProcessPlatform[_R] = Callable[[HomeAssistant, str, Any], _R | Awaitable[_R]] class LazyIntegrationPlatforms[_R]: @@ -277,6 +277,8 @@ class LazyIntegrationPlatforms[_R]: this only imports and processes the platform for an integration the first time it is requested, and only for integrations that are loaded. + The process callback may be a coroutine function; its result is awaited. + The platform is intentionally not registered for preloading, since for a rarely used platform that would import it for every integration during loading, defeating the point of loading it lazily. @@ -358,7 +360,10 @@ class LazyIntegrationPlatforms[_R]: result: _R | None = None if platform is not None: try: - result = self._process_platform(self._hass, domain, platform) + processed = self._process_platform(self._hass, domain, platform) + if isinstance(processed, Awaitable): + processed = await processed + result = processed except Exception: _LOGGER.exception( "Error processing %s platform for %s", diff --git a/tests/components/ai_task/test_media_source.py b/tests/components/ai_task/test_media_source.py index f41992f74cde..883f3953562b 100644 --- a/tests/components/ai_task/test_media_source.py +++ b/tests/components/ai_task/test_media_source.py @@ -26,6 +26,9 @@ async def test_local_media_source(hass: HomeAssistant, init_components: None) -> ) assert source.url_prefix == "/ai_task" + +async def test_media_source_no_media_dirs(hass: HomeAssistant) -> None: + """Test an error is raised when no media directories are configured.""" hass.config.media_dirs = {} with pytest.raises( diff --git a/tests/components/media_source/test_helper.py b/tests/components/media_source/test_helper.py index 03be63acc4bc..0f5a236ac6d4 100644 --- a/tests/components/media_source/test_helper.py +++ b/tests/components/media_source/test_helper.py @@ -1,6 +1,6 @@ """Test media source helpers.""" -from unittest.mock import Mock, patch +from unittest.mock import AsyncMock, patch import pytest @@ -11,7 +11,7 @@ from homeassistant.components.media_player import ( SearchMediaQuery, ) from homeassistant.components.media_source import const, models -from homeassistant.components.media_source.const import MEDIA_SOURCE_DATA +from homeassistant.components.media_source.const import DATA_MEDIA_SOURCE_PLATFORMS from homeassistant.core import HomeAssistant from homeassistant.setup import async_setup_component @@ -125,18 +125,18 @@ async def test_async_unresolve_media(hass: HomeAssistant) -> None: ) -async def test_browse_resolve_without_setup() -> None: +async def test_browse_resolve_without_setup(hass: HomeAssistant) -> None: """Test browse and resolve work without being setup.""" with pytest.raises(BrowseError): - await media_source.async_browse_media(Mock(data={}), None) + await media_source.async_browse_media(hass, None) with pytest.raises(BrowseError): await media_source.async_search_media( - Mock(data={}), None, SearchMediaQuery(search_query="test") + hass, None, SearchMediaQuery(search_query="test") ) with pytest.raises(media_source.Unresolvable): - await media_source.async_resolve_media(Mock(data={}), None, None) + await media_source.async_resolve_media(hass, None, None) async def test_async_search_media(hass: HomeAssistant) -> None: @@ -166,7 +166,11 @@ async def test_async_search_media(hass: HomeAssistant) -> None: async def test_async_search_media_not_supported(hass: HomeAssistant) -> None: """Test searching a source without search support raises a BrowseError.""" - hass.data[MEDIA_SOURCE_DATA] = {"plain": models.MediaSource("plain")} + assert await async_setup_component(hass, media_source.DOMAIN, {}) + await hass.async_block_till_done() + hass.data[DATA_MEDIA_SOURCE_PLATFORMS].async_get_platform = AsyncMock( + return_value=models.MediaSource("plain") + ) with pytest.raises(BrowseError): await media_source.async_search_media( @@ -178,10 +182,11 @@ async def test_async_search_media_not_supported(hass: HomeAssistant) -> None: async def test_async_search_media_root_not_supported(hass: HomeAssistant) -> None: """Test searching the aggregate root of multiple sources is not supported.""" - hass.data[MEDIA_SOURCE_DATA] = { - "source_a": models.MediaSource("source_a"), - "source_b": models.MediaSource("source_b"), - } + assert await async_setup_component(hass, media_source.DOMAIN, {}) + await hass.async_block_till_done() + hass.data[DATA_MEDIA_SOURCE_PLATFORMS].async_get_platforms = AsyncMock( + return_value={"source_a": models.MediaSource("source_a")} + ) with pytest.raises(BrowseError): await media_source.async_search_media( diff --git a/tests/helpers/test_integration_platform.py b/tests/helpers/test_integration_platform.py index 6784f1d9f083..0ad53bce417b 100644 --- a/tests/helpers/test_integration_platform.py +++ b/tests/helpers/test_integration_platform.py @@ -471,3 +471,25 @@ async def test_lazy_integration_platforms_concurrent(hass: HomeAssistant) -> Non assert results == [loaded_platform, loaded_platform] # The platform was imported and processed exactly once. assert processed == ["loaded"] + + +async def test_lazy_integration_platforms_async_process(hass: HomeAssistant) -> None: + """Test a coroutine process callback is awaited and its result cached.""" + loaded_platform = Mock() + mock_platform(hass, "loaded.platform_to_check", loaded_platform) + hass.config.components.add("loaded") + + processed: list[str] = [] + + async def _process_platform(hass: HomeAssistant, domain: str, platform: Any) -> Any: + processed.append(domain) + return platform + + platforms = LazyIntegrationPlatforms(hass, "platform_to_check", _process_platform) + + assert await platforms.async_get_platform("loaded") is loaded_platform + assert processed == ["loaded"] + + # The awaited result is cached, so a subsequent request does not reprocess. + assert await platforms.async_get_platform("loaded") is loaded_platform + assert processed == ["loaded"]