Added download types

1 files changed, 50 insertions, 0 deletions

diff --git a/src/model/youload_playlist.py b/src/model/youload_playlist.py
new file mode 100644
index 0000000..f23533f
--- /dev/null
+++ b/src/model/youload_playlist.py
@@ -0,0 +1,50 @@
+from pytube import Playlist
+import os
+
+def make_alpha_numeric(string: str) -> str:
+    return "".join(char for char in string if char.isalnum())
+
+class YouLoadPlayList:
+    """A class for download and handling youtube playlists"""
+
+    DOWNLOAD_TYPES = ["highest", "lowest", "audio", "720p", "480p", "360p", "240p", "144p"]
+    DEFAULT_DOWNLOAD_TYPE = "highest"
+
+    def __init__(self, link: str):
+        self.yt_playlist = Playlist(link)
+        self.video_count = len(self.yt_playlist.videos)
+        self.folder_name = make_alpha_numeric(self.yt_playlist.title)
+        self.download_type = self.DEFAULT_DOWNLOAD_TYPE
+
+    def generate_video_info(self, video_num: int) -> str:
+        """Returns information on video in playlist at 'video_num'"""
+
+        video = self.yt_playlist.videos[video_num]
+        video_size = video.streams.get_highest_resolution().filesize // 1048576
+        return f"Title: {video.title}, Size: {video_size} MB"
+
+    def set_download_directory(self, directory: str) -> None:
+        """Sets where the playlist folder will be downloaded"""
+        self.folder_name = os.path.join(directory, make_alpha_numeric(self.yt_playlist.title))
+
+    def prepare_for_download(self) -> None:
+        """Gets the playlist ready for download. Creates the output folder and that stuff."""
+        os.mkdir(self.folder_name)
+
+    def download_video(self, video_num: int) -> str:
+        """Download video at 'video_num'"""
+        video = self.yt_playlist.videos[video_num]
+        video_size = video.streams.get_highest_resolution().filesize
+
+        # Download this fucker.
+        if self.download_type == "highest":
+            video.streams.get_highest_resolution().download(output_path=self.folder_name)
+        elif self.download_type == "lowest":
+            video.streams.get_lowest_resolution().download(output_path=self.folder_name)
+        elif self.download_type == "audio":
+            video.streams.get_audio_only().download(output_path=self.folder_name)
+        else:
+            video.streams.get_by_resolution(self.download_type).download(output_path=self.folder_name)
+
+        return f"Title: {video.title}, Size: {video_size // (1024 ** 2)} MB"
+