From 56de00272704ead14c911da2325cad82920aa291 Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <balloob@gmail.com>
Date: Mon, 2 May 2022 06:09:49 -0700
Subject: [PATCH] Add media source support to AppleTV (#71185)

---
 .../components/apple_tv/browse_media.py       |  2 +-
 .../components/apple_tv/media_player.py       | 54 +++++++++++++++++--
 2 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/apple_tv/browse_media.py b/homeassistant/components/apple_tv/browse_media.py
index 3c0eee8b6ad..8d0a94ca858 100644
--- a/homeassistant/components/apple_tv/browse_media.py
+++ b/homeassistant/components/apple_tv/browse_media.py
@@ -18,7 +18,7 @@ def build_app_list(app_list):
 
     return BrowseMedia(
         media_class=MEDIA_CLASS_DIRECTORY,
-        media_content_id=None,
+        media_content_id="apps",
         media_content_type=MEDIA_TYPE_APPS,
         title="Apps",
         can_play=True,
diff --git a/homeassistant/components/apple_tv/media_player.py b/homeassistant/components/apple_tv/media_player.py
index cadf84e8b31..02919043bb2 100644
--- a/homeassistant/components/apple_tv/media_player.py
+++ b/homeassistant/components/apple_tv/media_player.py
@@ -13,11 +13,15 @@ from pyatv.const import (
 )
 from pyatv.helpers import is_streamable
 
+from homeassistant.components import media_source
 from homeassistant.components.media_player import (
     BrowseMedia,
     MediaPlayerEntity,
     MediaPlayerEntityFeature,
 )
+from homeassistant.components.media_player.browse_media import (
+    async_process_play_media_url,
+)
 from homeassistant.components.media_player.const import (
     MEDIA_TYPE_APP,
     MEDIA_TYPE_MUSIC,
@@ -73,7 +77,8 @@ SUPPORT_APPLE_TV = (
 
 # Map features in pyatv to Home Assistant
 SUPPORT_FEATURE_MAPPING = {
-    FeatureName.PlayUrl: MediaPlayerEntityFeature.PLAY_MEDIA,
+    FeatureName.PlayUrl: MediaPlayerEntityFeature.BROWSE_MEDIA
+    | MediaPlayerEntityFeature.PLAY_MEDIA,
     FeatureName.StreamFile: MediaPlayerEntityFeature.PLAY_MEDIA,
     FeatureName.Pause: MediaPlayerEntityFeature.PAUSE,
     FeatureName.Play: MediaPlayerEntityFeature.PLAY,
@@ -276,12 +281,24 @@ class AppleTvMediaPlayer(AppleTVEntity, MediaPlayerEntity):
         # RAOP. Otherwise try to play it with regular AirPlay.
         if media_type == MEDIA_TYPE_APP:
             await self.atv.apps.launch_app(media_id)
-        elif self._is_feature_available(FeatureName.StreamFile) and (
-            await is_streamable(media_id) or media_type == MEDIA_TYPE_MUSIC
+
+        is_media_source_id = media_source.is_media_source_id(media_id)
+
+        if (
+            not is_media_source_id
+            and self._is_feature_available(FeatureName.StreamFile)
+            and (await is_streamable(media_id) or media_type == MEDIA_TYPE_MUSIC)
         ):
             _LOGGER.debug("Streaming %s via RAOP", media_id)
             await self.atv.stream.stream_file(media_id)
-        elif self._is_feature_available(FeatureName.PlayUrl):
+
+        if self._is_feature_available(FeatureName.PlayUrl):
+            if is_media_source_id:
+                play_item = await media_source.async_resolve_media(self.hass, media_id)
+                media_id = play_item.url
+
+            media_id = async_process_play_media_url(self.hass, media_id)
+
             _LOGGER.debug("Playing %s via AirPlay", media_id)
             await self.atv.stream.play_url(media_id)
         else:
@@ -380,7 +397,34 @@ class AppleTvMediaPlayer(AppleTVEntity, MediaPlayerEntity):
         media_content_id=None,
     ) -> BrowseMedia:
         """Implement the websocket media browsing helper."""
-        return build_app_list(self._app_list)
+        # If we can't stream URLs, we can't browse media.
+        # In that case the `BROWSE_MEDIA` feature was added because of AppList/LaunchApp
+        if not self._is_feature_available(FeatureName.PlayUrl):
+            return build_app_list(self._app_list)
+
+        if self._app_list:
+            kwargs = {}
+        else:
+            # If it has no apps, assume it has no display
+            kwargs = {
+                "content_filter": lambda item: item.media_content_type.startswith(
+                    "audio/"
+                ),
+            }
+
+        cur_item = await media_source.async_browse_media(
+            self.hass, media_content_id, **kwargs
+        )
+
+        # If media content id is not None, we're browsing into a media source
+        if media_content_id is not None:
+            return cur_item
+
+        # Add app item if we have one
+        if self._app_list and cur_item.children:
+            cur_item.children.insert(0, build_app_list(self._app_list))
+
+        return cur_item
 
     async def async_turn_on(self):
         """Turn the media player on."""