[media_extractor] Add support for custom stream queries for media_extractor (#8538)

* Add support for different stream formats

* Encapsulate logic inside MediaExtractor class

* Add CONFIG_SCHEMA

* Fix for cases when youtube-dl returns content of playlist as list

homeassistant/components/media_extractor.py

@@ -6,11 +6,14 @@ https://home-assistant.io/components/media_extractor/
 """
 import logging
 import os
+import voluptuous as vol
 
 from homeassistant.components.media_player import (
-    ATTR_MEDIA_CONTENT_ID, DOMAIN as MEDIA_PLAYER_DOMAIN,
+    ATTR_ENTITY_ID, ATTR_MEDIA_CONTENT_ID, ATTR_MEDIA_CONTENT_TYPE,
-    MEDIA_PLAYER_PLAY_MEDIA_SCHEMA, SERVICE_PLAY_MEDIA)
+    DOMAIN as MEDIA_PLAYER_DOMAIN, MEDIA_PLAYER_PLAY_MEDIA_SCHEMA,
+    SERVICE_PLAY_MEDIA)
 from homeassistant.config import load_yaml_config_file
+from homeassistant.helpers import config_validation as cv
 
 REQUIREMENTS = ['youtube_dl==2017.7.9']
 
@@ -19,6 +22,18 @@ _LOGGER = logging.getLogger(__name__)
 DOMAIN = 'media_extractor'
 DEPENDENCIES = ['media_player']
 
+CONF_CUSTOMIZE_ENTITIES = 'customize'
+CONF_DEFAULT_STREAM_QUERY = 'default_query'
+DEFAULT_STREAM_QUERY = 'best'
+
+CONFIG_SCHEMA = vol.Schema({
+    DOMAIN: vol.Schema({
+        vol.Optional(CONF_DEFAULT_STREAM_QUERY): cv.string,
+        vol.Optional(CONF_CUSTOMIZE_ENTITIES):
+            vol.Schema({cv.entity_id: vol.Schema({cv.string: cv.string})}),
+    }),
+}, extra=vol.ALLOW_EXTRA)
+
 
 def setup(hass, config):
     """Set up the media extractor service."""
@@ -28,23 +43,7 @@ def setup(hass, config):
 
     def play_media(call):
         """Get stream URL and send it to the media_player.play_media."""
-        media_url = call.data.get(ATTR_MEDIA_CONTENT_ID)
+        MediaExtractor(hass, config[DOMAIN], call.data).extract_and_send()
-
-        try:
-            stream_url = get_media_stream_url(media_url)
-        except YDException:
-            _LOGGER.error("Could not retrieve data for the URL: %s",
-                          media_url)
-            return
-        else:
-            data = {k: v for k, v in call.data.items()
-                    if k != ATTR_MEDIA_CONTENT_ID}
-            data[ATTR_MEDIA_CONTENT_ID] = stream_url
-
-            hass.async_add_job(
-                hass.services.async_call(
-                    MEDIA_PLAYER_DOMAIN, SERVICE_PLAY_MEDIA, data)
-            )
 
     hass.services.register(DOMAIN,
                            SERVICE_PLAY_MEDIA,
@@ -55,47 +54,136 @@ def setup(hass, config):
     return True
 
 
-class YDException(Exception):
+class MEDownloadException(Exception):
-    """General service exception."""
+    """Media extractor download exception."""
 
     pass
 
 
-def get_media_stream_url(media_url):
+class MEQueryException(Exception):
-    """Extract stream URL from the media URL."""
+    """Media extractor query exception."""
+
+    pass
+
+
+class MediaExtractor:
+    """Class which encapsulates all extraction logic."""
+
+    def __init__(self, hass, component_config, call_data):
+        """Initialize media extractor."""
+        self.hass = hass
+        self.config = component_config
+        self.call_data = call_data
+
+    def get_media_url(self):
+        """Return media content url."""
+        return self.call_data.get(ATTR_MEDIA_CONTENT_ID)
+
+    def get_entities(self):
+        """Return list of entities."""
+        return self.call_data.get(ATTR_ENTITY_ID, [])
+
+    def extract_and_send(self):
+        """Extract exact stream format for each entity_id and play it."""
+        try:
+            stream_selector = self.get_stream_selector()
+        except MEDownloadException:
+            _LOGGER.error("Could not retrieve data for the URL: %s",
+                          self.get_media_url())
+        else:
+            entities = self.get_entities()
+
+            if len(entities) == 0:
+                self.call_media_player_service(stream_selector, None)
+
+            for entity_id in entities:
+                self.call_media_player_service(stream_selector, entity_id)
+
+    def get_stream_selector(self):
+        """Return format selector for the media URL."""
         from youtube_dl import YoutubeDL
         from youtube_dl.utils import DownloadError, ExtractorError
 
         ydl = YoutubeDL({'quiet': True, 'logger': _LOGGER})
 
         try:
-        all_media_streams = ydl.extract_info(media_url, process=False)
+            all_media = ydl.extract_info(self.get_media_url(),
+                                         process=False)
         except DownloadError:
             # This exception will be logged by youtube-dl itself
-        raise YDException()
+            raise MEDownloadException()
 
-    if 'entries' in all_media_streams:
+        if 'entries' in all_media:
             _LOGGER.warning("Playlists are not supported, "
                             "looking for the first video")
-        try:
+            entries = list(all_media['entries'])
-            selected_stream = next(all_media_streams['entries'])
+            if len(entries) > 0:
-        except StopIteration:
+                selected_media = entries[0]
-            _LOGGER.error("Playlist is empty")
-            raise YDException()
             else:
-        selected_stream = all_media_streams
+                _LOGGER.error("Playlist is empty")
+                raise MEDownloadException()
+        else:
+            selected_media = all_media
 
         try:
-        media_info = ydl.process_ie_result(selected_stream, download=False)
+            media_info = ydl.process_ie_result(selected_media,
+                                               download=False)
         except (ExtractorError, DownloadError):
             # This exception will be logged by youtube-dl itself
-        raise YDException()
+            raise MEDownloadException()
 
-    format_selector = ydl.build_format_selector('best')
+        def stream_selector(query):
+            """Find stream url that matches query."""
+            try:
+                format_selector = ydl.build_format_selector(query)
+            except (SyntaxError, ValueError, AttributeError) as ex:
+                _LOGGER.error(ex)
+                raise MEQueryException()
 
             try:
-        best_quality_stream = next(format_selector(media_info))
+                requested_stream = next(format_selector(media_info))
             except (KeyError, StopIteration):
-        best_quality_stream = media_info
+                _LOGGER.error("Could not extract stream for the query: %s",
+                              query)
+                raise MEQueryException()
 
-    return best_quality_stream['url']
+            return requested_stream['url']
+
+        return stream_selector
+
+    def call_media_player_service(self, stream_selector, entity_id):
+        """Call media_player.play_media service."""
+        stream_query = self.get_stream_query_for_entity(entity_id)
+
+        try:
+            stream_url = stream_selector(stream_query)
+        except MEQueryException:
+            _LOGGER.error("Wrong query format: %s", stream_query)
+            return
+        else:
+            data = {k: v for k, v in self.call_data.items()
+                    if k != ATTR_ENTITY_ID}
+            data[ATTR_MEDIA_CONTENT_ID] = stream_url
+
+            if entity_id:
+                data[ATTR_ENTITY_ID] = entity_id
+
+            self.hass.async_add_job(
+                self.hass.services.async_call(
+                    MEDIA_PLAYER_DOMAIN, SERVICE_PLAY_MEDIA, data)
+            )
+
+    def get_stream_query_for_entity(self, entity_id):
+        """Get stream format query for entity."""
+        default_stream_query = self.config.get(CONF_DEFAULT_STREAM_QUERY,
+                                               DEFAULT_STREAM_QUERY)
+
+        if entity_id:
+            media_content_type = self.call_data.get(ATTR_MEDIA_CONTENT_TYPE)
+
+            return self.config \
+                .get(CONF_CUSTOMIZE_ENTITIES, {}) \
+                .get(entity_id, {}) \
+                .get(media_content_type, default_stream_query)
+
+        return default_stream_query