From 655b57d2a871e11722cfa4b24651333a7a738f45 Mon Sep 17 00:00:00 2001
From: Maksym Pavlenko <pavlenko.maksym@gmail.com>
Date: Sun, 14 Apr 2019 18:16:25 -0700
Subject: [PATCH] Fix resolver for audio (return m4a file instead of XML)

---
 cmd/resolver/function.py      | 16 +++++++++++++---
 cmd/resolver/function_test.py |  6 ++++++
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/cmd/resolver/function.py b/cmd/resolver/function.py
index 12af731..0dd301b 100644
--- a/cmd/resolver/function.py
+++ b/cmd/resolver/function.py
@@ -137,7 +137,7 @@ def _resolve(url, metadata):
         with youtube_dl.YoutubeDL(opts) as ytdl:
             info = ytdl.extract_info(url, download=False)
             if provider == 'youtube':
-                return _yt_choose_url(info, metadata)
+                return _yt_choose_url(ytdl, info, metadata)
             elif provider == 'vimeo':
                 return _vimeo_choose_url(info, metadata)
             else:
@@ -147,8 +147,19 @@ def _resolve(url, metadata):
         raise
 
 
-def _yt_choose_url(info, metadata):
+def _yt_choose_url(ytdl, info, metadata):
     is_video = metadata['format'] == 'video'
+    is_high_quality = metadata['quality'] == 'high'
+
+    if not is_video:
+        fmt = 'bestaudio' if is_high_quality else 'worstaudio'
+        selector = ytdl.build_format_selector(fmt)
+        try:
+            selected = next(selector(info))
+            if 'fragment_base_url' in selected:
+                return selected['fragment_base_url']
+        except KeyError:
+            pass
 
     # Filter formats by file extension
     ext = 'mp4' if is_video else 'm4a'
@@ -164,7 +175,6 @@ def _yt_choose_url(info, metadata):
     ordered = sorted(fmt_list, key=lambda x: x[sort_field], reverse=True)
 
     # Choose an item depending on quality, better at the beginning
-    is_high_quality = metadata['quality'] == 'high'
     item = ordered[0] if is_high_quality else ordered[-1]
     return item['url']
 
diff --git a/cmd/resolver/function_test.py b/cmd/resolver/function_test.py
index 8793682..ea4b464 100644
--- a/cmd/resolver/function_test.py
+++ b/cmd/resolver/function_test.py
@@ -26,6 +26,12 @@ class TestYtdl(unittest.TestCase):
             ytdl._resolve('https://vimeo.com/275211960', {'format': 'video', 'quality': 'high', 'provider': 'vimeo'})
         )
 
+    def test_youtube_resolve_audio(self):
+        self.assertIsNotNone(
+            ytdl._resolve('https://youtube.com/watch?v=UMrb1tG38w8',
+                          {'format': 'audio', 'quality': 'high', 'provider': 'youtube'})
+        )
+
 
 class TestDynamo(unittest.TestCase):
     def test_metadata(self):