From eba1ec11c78645c2bcfbb7ec04f2df4faf4cd1cf Mon Sep 17 00:00:00 2001
From: shirt <2660574+shirt-dev@users.noreply.github.com>
Date: Tue, 26 Oct 2021 15:06:17 -0400
Subject: [PATCH] Use ffmpeg instead of pydub, support ogg stream copy

---
 requirements.txt  |  6 +++---
 zspotify/const.py | 23 +++++++++++++++++++++++
 zspotify/track.py | 46 ++++++++++++++++++++++++++--------------------
 3 files changed, 52 insertions(+), 23 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 4f4b877..7163d55 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,6 @@
+ffmpy
 git+https://github.com/kokarare1212/librespot-python
 music_tag
-pydub
 Pillow
-tqdm
-tabulate
\ No newline at end of file
+tabulate
+tqdm
\ No newline at end of file
diff --git a/zspotify/const.py b/zspotify/const.py
index 3e4d9b0..3708c58 100644
--- a/zspotify/const.py
+++ b/zspotify/const.py
@@ -96,11 +96,34 @@ CHUNK_SIZE = 'CHUNK_SIZE'
 
 SPLIT_ALBUM_DISCS = 'SPLIT_ALBUM_DISCS'
 
+BITRATE = 'BITRATE'
+
+CODEC_MAP = {
+    'aac': 'aac',
+    'fdk_aac': 'libfdk_aac',
+    'm4a': 'aac',
+    'mp3': 'libmp3lame',
+    'ogg': 'copy',
+    'opus': 'libopus',
+    'vorbis': 'copy',
+}
+
+EXT_MAP = {
+    'aac': 'm4a',
+    'fdk_aac': 'm4a',
+    'm4a': 'm4a',
+    'mp3': 'mp3',
+    'ogg': 'ogg',
+    'opus': 'ogg',
+    'vorbis': 'ogg',
+}
+
 CONFIG_DEFAULT_SETTINGS = {
     'ROOT_PATH': '../ZSpotify Music/',
     'ROOT_PODCAST_PATH': '../ZSpotify Podcasts/',
     'SKIP_EXISTING_FILES': True,
     'DOWNLOAD_FORMAT': 'mp3',
+    'BITRATE': '160k',
     'FORCE_PREMIUM': False,
     'ANTI_BAN_WAIT_TIME': 1,
     'OVERRIDE_AUTO_WAIT': False,
diff --git a/zspotify/track.py b/zspotify/track.py
index 423cce0..87452de 100644
--- a/zspotify/track.py
+++ b/zspotify/track.py
@@ -2,16 +2,14 @@ import os
 import time
 from typing import Any, Tuple, List
 
-from librespot.audio.decoders import AudioQuality
 from librespot.metadata import TrackId
-from pydub import AudioSegment
+from ffmpy import FFmpeg
 from tqdm import tqdm
 
 from const import TRACKS, ALBUM, NAME, ITEMS, DISC_NUMBER, TRACK_NUMBER, IS_PLAYABLE, ARTISTS, IMAGES, URL, \
     RELEASE_DATE, ID, TRACKS_URL, SAVED_TRACKS_URL, SPLIT_ALBUM_DISCS, ROOT_PATH, DOWNLOAD_FORMAT, CHUNK_SIZE, \
-    SKIP_EXISTING_FILES, ANTI_BAN_WAIT_TIME, OVERRIDE_AUTO_WAIT
-from utils import sanitize_data, set_audio_tags, set_music_thumbnail, create_download_directory, \
-    MusicFormat
+    SKIP_EXISTING_FILES, ANTI_BAN_WAIT_TIME, OVERRIDE_AUTO_WAIT, BITRATE, CODEC_MAP, EXT_MAP
+from utils import sanitize_data, set_audio_tags, set_music_thumbnail, create_download_directory
 from zspotify import ZSpotify
 
 
@@ -72,7 +70,7 @@ def download_track(track_id: str, extra_paths='', prefix=False, prefix_value='',
             ) else f'{prefix_value} - {song_name}'
 
         filename = os.path.join(
-            download_directory, f'{song_name}.{ZSpotify.get_config(DOWNLOAD_FORMAT)}')
+            download_directory, f'{song_name}.{EXT_MAP.get(ZSpotify.get_config(DOWNLOAD_FORMAT))}')
 
     except Exception as e:
         print('###   SKIPPING SONG - FAILED TO QUERY METADATA   ###')
@@ -107,11 +105,10 @@ def download_track(track_id: str, extra_paths='', prefix=False, prefix_value='',
                             p_bar.update(file.write(
                                 stream.input_stream.stream().read(ZSpotify.get_config(CHUNK_SIZE))))
 
-                    if ZSpotify.get_config(DOWNLOAD_FORMAT) == 'mp3':
-                        convert_audio_format(filename)
-                        set_audio_tags(filename, artists, name, album_name,
-                                       release_year, disc_number, track_number)
-                        set_music_thumbnail(filename, image_url)
+                    convert_audio_format(filename)
+                    set_audio_tags(filename, artists, name, album_name,
+                                   release_year, disc_number, track_number)
+                    set_music_thumbnail(filename, image_url)
 
                     if not ZSpotify.get_config(OVERRIDE_AUTO_WAIT):
                         time.sleep(ZSpotify.get_config(ANTI_BAN_WAIT_TIME))
@@ -124,13 +121,22 @@ def download_track(track_id: str, extra_paths='', prefix=False, prefix_value='',
 
 
 def convert_audio_format(filename) -> None:
-    """ Converts raw audio into playable mp3 """
-    # print('###   CONVERTING TO ' + MUSIC_FORMAT.upper() + '   ###')
-    raw_audio = AudioSegment.from_file(filename, format=MusicFormat.OGG.value,
-                                       frame_rate=44100, channels=2, sample_width=2)
-    if ZSpotify.DOWNLOAD_QUALITY == AudioQuality.VERY_HIGH:
-        bitrate = '320k'
+    """ Converts raw audio into playable file """
+    temp_filename = f'{os.path.splitext(filename)[0]}.tmp'
+    os.replace(filename, temp_filename)
+
+    download_format = ZSpotify.get_config(DOWNLOAD_FORMAT)
+    file_codec = CODEC_MAP.get(download_format, "copy")
+    if file_codec != 'copy':
+        bitrate = ZSpotify.get_config(BITRATE)
     else:
-        bitrate = '160k'
-    raw_audio.export(filename, format=ZSpotify.get_config(
-        DOWNLOAD_FORMAT), bitrate=bitrate)
+        bitrate = None
+
+    ff_m = FFmpeg(
+        global_options=['-y', '-hide_banner', '-loglevel error'],
+        inputs={temp_filename: None},
+        outputs={filename: ['-c:a', file_codec] + ['-b:a', bitrate] if bitrate else []}
+    )
+    ff_m.run()
+    if os.path.exists(temp_filename):
+        os.remove(temp_filename)