]> jfr.im git - yt-dlp.git/blobdiff - test/test_YoutubeDL.py
Add slicing notation to `--playlist-items`
[yt-dlp.git] / test / test_YoutubeDL.py
index 1133f61654cf922d8704aa5a3091fa638fcecd74..3aafc3c4f7215e5123aa62af9368407b374506ad 100644 (file)
@@ -23,6 +23,7 @@
 from yt_dlp.utils import (
     ExtractorError,
     LazyList,
+    OnDemandPagedList,
     int_or_none,
     match_filter_func,
 )
@@ -989,41 +990,79 @@ def f(v, incomplete):
         self.assertEqual(res, [])
 
     def test_playlist_items_selection(self):
-        entries = [{
-            'id': compat_str(i),
-            'title': compat_str(i),
-            'url': TEST_URL,
-        } for i in range(1, 5)]
-        playlist = {
-            '_type': 'playlist',
-            'id': 'test',
-            'entries': entries,
-            'extractor': 'test:playlist',
-            'extractor_key': 'test:playlist',
-            'webpage_url': 'http://example.com',
-        }
+        INDICES, PAGE_SIZE = list(range(1, 11)), 3
+
+        def entry(i, evaluated):
+            evaluated.append(i)
+            return {
+                'id': str(i),
+                'title': str(i),
+                'url': TEST_URL,
+            }
+
+        def pagedlist_entries(evaluated):
+            def page_func(n):
+                start = PAGE_SIZE * n
+                for i in INDICES[start: start + PAGE_SIZE]:
+                    yield entry(i, evaluated)
+            return OnDemandPagedList(page_func, PAGE_SIZE)
+
+        def page_num(i):
+            return (i + PAGE_SIZE - 1) // PAGE_SIZE
 
-        def get_downloaded_info_dicts(params):
+        def generator_entries(evaluated):
+            for i in INDICES:
+                yield entry(i, evaluated)
+
+        def list_entries(evaluated):
+            return list(generator_entries(evaluated))
+
+        def lazylist_entries(evaluated):
+            return LazyList(generator_entries(evaluated))
+
+        def get_downloaded_info_dicts(params, entries):
             ydl = YDL(params)
-            # make a deep copy because the dictionary and nested entries
-            # can be modified
-            ydl.process_ie_result(copy.deepcopy(playlist))
+            ydl.process_ie_result({
+                '_type': 'playlist',
+                'id': 'test',
+                'extractor': 'test:playlist',
+                'extractor_key': 'test:playlist',
+                'webpage_url': 'http://example.com',
+                'entries': entries,
+            })
             return ydl.downloaded_info_dicts
 
-        def test_selection(params, expected_ids):
-            results = [
-                (v['playlist_autonumber'] - 1, (int(v['id']), v['playlist_index']))
-                for v in get_downloaded_info_dicts(params)]
-            self.assertEqual(results, list(enumerate(zip(expected_ids, expected_ids))))
-
-        test_selection({}, [1, 2, 3, 4])
-        test_selection({'playlistend': 10}, [1, 2, 3, 4])
-        test_selection({'playlistend': 2}, [1, 2])
-        test_selection({'playliststart': 10}, [])
-        test_selection({'playliststart': 2}, [2, 3, 4])
-        test_selection({'playlist_items': '2-4'}, [2, 3, 4])
+        def test_selection(params, expected_ids, evaluate_all=False):
+            expected_ids = list(expected_ids)
+            if evaluate_all:
+                generator_eval = pagedlist_eval = INDICES
+            elif not expected_ids:
+                generator_eval = pagedlist_eval = []
+            else:
+                generator_eval = INDICES[0: max(expected_ids)]
+                pagedlist_eval = INDICES[PAGE_SIZE * page_num(min(expected_ids)) - PAGE_SIZE:
+                                         PAGE_SIZE * page_num(max(expected_ids))]
+
+            for name, func, expected_eval in (
+                ('list', list_entries, INDICES),
+                ('Generator', generator_entries, generator_eval),
+                ('LazyList', lazylist_entries, generator_eval),
+                ('PagedList', pagedlist_entries, pagedlist_eval),
+            ):
+                evaluated = []
+                entries = func(evaluated)
+                results = [(v['playlist_autonumber'] - 1, (int(v['id']), v['playlist_index']))
+                           for v in get_downloaded_info_dicts(params, entries)]
+                self.assertEqual(results, list(enumerate(zip(expected_ids, expected_ids))), f'Entries of {name} for {params}')
+                self.assertEqual(sorted(evaluated), expected_eval, f'Evaluation of {name} for {params}')
+        test_selection({}, INDICES)
+        test_selection({'playlistend': 20}, INDICES, True)
+        test_selection({'playlistend': 2}, INDICES[:2])
+        test_selection({'playliststart': 11}, [], True)
+        test_selection({'playliststart': 2}, INDICES[1:])
+        test_selection({'playlist_items': '2-4'}, INDICES[1:4])
         test_selection({'playlist_items': '2,4'}, [2, 4])
-        test_selection({'playlist_items': '10'}, [])
+        test_selection({'playlist_items': '20'}, [], True)
         test_selection({'playlist_items': '0'}, [])
 
         # Tests for https://github.com/ytdl-org/youtube-dl/issues/10591
@@ -1032,11 +1071,33 @@ def test_selection(params, expected_ids):
 
         # Tests for https://github.com/yt-dlp/yt-dlp/issues/720
         # https://github.com/yt-dlp/yt-dlp/issues/302
-        test_selection({'playlistreverse': True}, [4, 3, 2, 1])
-        test_selection({'playliststart': 2, 'playlistreverse': True}, [4, 3, 2])
+        test_selection({'playlistreverse': True}, INDICES[::-1])
+        test_selection({'playliststart': 2, 'playlistreverse': True}, INDICES[:0:-1])
         test_selection({'playlist_items': '2,4', 'playlistreverse': True}, [4, 2])
         test_selection({'playlist_items': '4,2'}, [4, 2])
 
+        # Tests for --playlist-items start:end:step
+        test_selection({'playlist_items': ':'}, INDICES, True)
+        test_selection({'playlist_items': '::1'}, INDICES, True)
+        test_selection({'playlist_items': '::-1'}, INDICES[::-1], True)
+        test_selection({'playlist_items': ':6'}, INDICES[:6])
+        test_selection({'playlist_items': ':-6'}, INDICES[:-5], True)
+        test_selection({'playlist_items': '-1:6:-2'}, INDICES[:4:-2], True)
+        test_selection({'playlist_items': '9:-6:-2'}, INDICES[8:3:-2], True)
+
+        test_selection({'playlist_items': '1:inf:2'}, INDICES[::2], True)
+        test_selection({'playlist_items': '-2:inf'}, INDICES[-2:], True)
+        test_selection({'playlist_items': ':inf:-1'}, [], True)
+        test_selection({'playlist_items': '0-2:2'}, [2])
+        test_selection({'playlist_items': '1-:2'}, INDICES[::2], True)
+        test_selection({'playlist_items': '0--2:2'}, INDICES[1:-1:2], True)
+
+        test_selection({'playlist_items': '10::3'}, [10], True)
+        test_selection({'playlist_items': '-1::3'}, [10], True)
+        test_selection({'playlist_items': '11::3'}, [], True)
+        test_selection({'playlist_items': '-15::2'}, INDICES[1::2], True)
+        test_selection({'playlist_items': '-15::15'}, [], True)
+
     def test_urlopen_no_file_protocol(self):
         # see https://github.com/ytdl-org/youtube-dl/issues/8227
         ydl = YDL()