Merge pull request #253 from johnvanderholt/master

pytube/contrib/playlist.py

@@ -66,26 +66,65 @@ class Playlist(object):
             complete_url = base_url + video_id
             self.video_urls.append(complete_url)
 
-    def download_all(self, download_path=None):
+    def _path_num_prefix_generator(self, reverse=False):
+        """
+        This generator function generates number prefixes, for the items
+        in the playlist.
+        If the number of digits required to name a file,is less than is
+        required to name the last file,it prepends 0s.
+        So if you have a playlist of 100 videos it will number them like:
+        001, 002, 003 ect, up to 100.
+        It also adds a space after the number.
+        :return: prefix string generator : generator
+        """
+        digits = len(str(len(self.video_urls)))
+        if reverse:
+            start, stop, step = (len(self.video_urls), 0, -1)
+        else:
+            start, stop, step = (1, len(self.video_urls) + 1, 1)
+        return (str(i).zfill(digits) for i in range(start, stop, step))
+
+    def download_all(self, download_path=None, prefix_number=True,
+                     reverse_numbering=False):
         """Download all the videos in the the playlist. Initially, download
         resolution is 720p (or highest available), later more option
         should be added to download resolution of choice
 
         TODO(nficano): Add option to download resolution of user's choice
+
+        :param download_path:
+            (optional) Output path for the playlist If one is not
+            specified, defaults to the current working directory.
+            This is passed along to the Stream objects.
+        :type download_path: str or None
+        :param prefix_number:
+            (optional) Automatically numbers playlists using the
+            _path_num_prefix_generator function.
+        :type prefix_number: bool
+        :param reverse_numbering:
+            (optional) Lets you number playlists in reverse, since some
+            playlists are ordered newest -> oldests.
+        :type reverse_numbering: bool
         """
 
         self.populate_video_urls()
         logger.debug('total videos found: ', len(self.video_urls))
         logger.debug('starting download')
 
+        prefix_gen = self._path_num_prefix_generator(reverse_numbering)
+
         for link in self.video_urls:
             yt = YouTube(link)
-
             # TODO: this should not be hardcoded to a single user's preference
             dl_stream = yt.streams.filter(
                 progressive=True, subtype='mp4',
             ).order_by('resolution').desc().first()
 
             logger.debug('download path: %s', download_path)
-            dl_stream.download(download_path)
+            if prefix_number:
+                prefix = next(prefix_gen)
+                logger.debug('file prefix is: %s', prefix)
+                dl_stream.download(download_path, filename_prefix=prefix)
+            else:
+                dl_stream.download(download_path)
             logger.debug('download complete')

pytube/streams.py

@@ -174,7 +174,7 @@ class Stream(object):
         filename = safe_filename(title)
         return '{filename}.{s.subtype}'.format(filename=filename, s=self)
 
-    def download(self, output_path=None, filename=None):
+    def download(self, output_path=None, filename=None, filename_prefix=None):
         """Write the media stream to disk.
 
         :param output_path:
@@ -185,6 +185,13 @@ class Stream(object):
             (optional) Output filename (stem only) for writing media file.
             If one is not specified, the default filename is used.
         :type filename: str or None
+        :param filename_prefix:
+            (optional) A string that will be prepended to the filename.
+            For example a number in a playlist or the name of a series.
+            If one is not specified, nothing will be prepended
+            This is seperate from filename so you can use the default
+            filename but still add a prefix.
+        :type filename_prefix: str or None
 
         :rtype: str
 
@@ -195,6 +202,11 @@ class Stream(object):
             filename = '{filename}.{s.subtype}'.format(filename=safe, s=self)
         filename = filename or self.default_filename
 
+        if filename_prefix:
+            filename = "{prefix}{filename}"\
+                .format(prefix=safe_filename(filename_prefix),
+                        filename=filename)
+
         # file path
         fp = os.path.join(output_path, filename)
         bytes_remaining = self.filesize

tests/test_playlist.py

@@ -1,12 +1,14 @@
 # -*- coding: utf-8 -*-
 from pytube import Playlist
 
+short_test_pl = 'https://www.youtube.com/watch?v=' \
+          'm5q2GCsteQs&list=PL525f8ds9RvsXDl44X6Wwh9t3fCzFNApw'
+long_test_pl = "https://www.youtube.com/watch?v=" \
+               "9CHDoAsX1yo&list=UUXuqSBlHAE6Xw-yeJA0Tunw"
+
 
 def test_construct():
-    ob = Playlist(
-        'https://www.youtube.com/watch?v=m5q2GCsteQs&list='
-        'PL525f8ds9RvsXDl44X6Wwh9t3fCzFNApw',
-    )
+    ob = Playlist(short_test_pl)
     expected = 'https://www.youtube.com/' \
                'playlist?list=' \
                'PL525f8ds9RvsXDl44X6Wwh9t3fCzFNApw'
@@ -14,25 +16,8 @@ def test_construct():
     assert ob.construct_playlist_url() == expected
 
 
-def test_link_parse():
-    ob = Playlist(
-        'https://www.youtube.com/watch?v=m5q2GCsteQs&list='
-        'PL525f8ds9RvsXDl44X6Wwh9t3fCzFNApw',
-    )
-
-    expected = [
-        '/watch?v=m5q2GCsteQs',
-        '/watch?v=5YK63cXyJ2Q',
-        '/watch?v=Rzt4rUPFYD4',
-    ]
-    assert ob.parse_links() == expected
-
-
 def test_populate():
-    ob = Playlist(
-        'https://www.youtube.com/watch?v=m5q2GCsteQs&list='
-        'PL525f8ds9RvsXDl44X6Wwh9t3fCzFNApw',
-    )
+    ob = Playlist(short_test_pl)
     expected = [
         'https://www.youtube.com/watch?v=m5q2GCsteQs',
         'https://www.youtube.com/watch?v=5YK63cXyJ2Q',
@@ -43,10 +28,50 @@ def test_populate():
     assert ob.video_urls == expected
 
 
+def test_link_parse():
+    ob = Playlist(short_test_pl)
+
+    expected = [
+        '/watch?v=m5q2GCsteQs',
+        '/watch?v=5YK63cXyJ2Q',
+        '/watch?v=Rzt4rUPFYD4',
+    ]
+    assert ob.parse_links() == expected
+
+
 def test_download():
-    ob = Playlist(
-        'https://www.youtube.com/watch?v=lByG_AgKS9k&list='
-        'PL525f8ds9RvuerPZ3bZygmNiYw2sP4BDk',
-    )
+    ob = Playlist(short_test_pl)
     ob.download_all()
     ob.download_all('.')
+    ob.download_all(prefix_number=False)
+    ob.download_all('.', prefix_number=False)
+
+
+def test_numbering():
+    ob = Playlist(short_test_pl)
+    ob.populate_video_urls()
+    gen = ob._path_num_prefix_generator(reverse=False)
+    assert "1" in next(gen)
+    assert "2" in next(gen)
+
+    ob = Playlist(short_test_pl)
+    ob.populate_video_urls()
+    gen = ob._path_num_prefix_generator(reverse=True)
+    assert str(len(ob.video_urls)) in next(gen)
+    assert str(len(ob.video_urls) - 1) in next(gen)
+
+    ob = Playlist(long_test_pl)
+    ob.populate_video_urls()
+    gen = ob._path_num_prefix_generator(reverse=False)
+    nxt = next(gen)
+    assert len(nxt) > 1
+    assert "1" in nxt
+    nxt = next(gen)
+    assert len(nxt) > 1
+    assert "2" in nxt
+
+    ob = Playlist(long_test_pl)
+    ob.populate_video_urls()
+    gen = ob._path_num_prefix_generator(reverse=True)
+    assert str(len(ob.video_urls)) in next(gen)
+    assert str(len(ob.video_urls) - 1) in next(gen)