diff options
| -rw-r--r-- | test/test_utils.py | 43 | ||||
| -rw-r--r-- | yt_dlp/utils.py | 43 |
2 files changed, 77 insertions, 9 deletions
diff --git a/test/test_utils.py b/test/test_utils.py index 9ff13a369..ade10a7b1 100644 --- a/test/test_utils.py +++ b/test/test_utils.py @@ -12,6 +12,7 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) # Various small unit tests import io +import itertools import json import xml.etree.ElementTree @@ -108,6 +109,7 @@ from yt_dlp.utils import ( cli_bool_option, parse_codecs, iri_to_uri, + LazyList, ) from yt_dlp.compat import ( compat_chr, @@ -1525,6 +1527,47 @@ Line 1 self.assertEqual(clean_podcast_url('https://www.podtrac.com/pts/redirect.mp3/chtbl.com/track/5899E/traffic.megaphone.fm/HSW7835899191.mp3'), 'https://traffic.megaphone.fm/HSW7835899191.mp3') self.assertEqual(clean_podcast_url('https://play.podtrac.com/npr-344098539/edge1.pod.npr.org/anon.npr-podcasts/podcast/npr/waitwait/2020/10/20201003_waitwait_wwdtmpodcast201003-015621a5-f035-4eca-a9a1-7c118d90bc3c.mp3'), 'https://edge1.pod.npr.org/anon.npr-podcasts/podcast/npr/waitwait/2020/10/20201003_waitwait_wwdtmpodcast201003-015621a5-f035-4eca-a9a1-7c118d90bc3c.mp3') + def test_LazyList(self): + it = list(range(10)) + + self.assertEqual(list(LazyList(it)), it) + self.assertEqual(LazyList(it).exhaust(), it) + self.assertEqual(LazyList(it)[5], it[5]) + + self.assertEqual(LazyList(it)[::2], it[::2]) + self.assertEqual(LazyList(it)[1::2], it[1::2]) + self.assertEqual(LazyList(it)[6:2:-2], it[6:2:-2]) + self.assertEqual(LazyList(it)[::-1], it[::-1]) + + self.assertTrue(LazyList(it)) + self.assertFalse(LazyList(range(0))) + self.assertEqual(len(LazyList(it)), len(it)) + self.assertEqual(repr(LazyList(it)), repr(it)) + self.assertEqual(str(LazyList(it)), str(it)) + + self.assertEqual(list(reversed(LazyList(it))), it[::-1]) + self.assertEqual(list(reversed(LazyList(it))[1:3:7]), it[::-1][1:3:7]) + + def test_LazyList_laziness(self): + + def test(ll, idx, val, cache): + self.assertEqual(ll[idx], val) + self.assertEqual(getattr(ll, '_LazyList__cache'), list(cache)) + + ll = LazyList(range(10)) + test(ll, 0, 0, range(1)) + test(ll, 5, 5, range(6)) + test(ll, -3, 7, range(10)) + + ll = reversed(LazyList(range(10))) + test(ll, -1, 0, range(1)) + test(ll, 3, 6, range(10)) + + ll = LazyList(itertools.count()) + test(ll, 10, 10, range(11)) + reversed(ll) + test(ll, -15, 14, range(15)) + if __name__ == '__main__': unittest.main() diff --git a/yt_dlp/utils.py b/yt_dlp/utils.py index 96f41ddd4..3cb79b657 100644 --- a/yt_dlp/utils.py +++ b/yt_dlp/utils.py @@ -3954,10 +3954,14 @@ class LazyList(collections.Sequence): def __init__(self, iterable): self.__iterable = iter(iterable) self.__cache = [] + self.__reversed = False def __iter__(self): - for item in self.__cache: - yield item + if self.__reversed: + # We need to consume the entire iterable to iterate in reverse + yield from self.exhaust()[::-1] + return + yield from self.__cache for item in self.__iterable: self.__cache.append(item) yield item @@ -3965,29 +3969,39 @@ class LazyList(collections.Sequence): def exhaust(self): ''' Evaluate the entire iterable ''' self.__cache.extend(self.__iterable) + return self.__cache + + @staticmethod + def _reverse_index(x): + return -(x + 1) def __getitem__(self, idx): if isinstance(idx, slice): step = idx.step or 1 - start = idx.start if idx.start is not None else 1 if step > 0 else -1 + start = idx.start if idx.start is not None else 0 if step > 0 else -1 stop = idx.stop if idx.stop is not None else -1 if step > 0 else 0 + if self.__reversed: + start, stop, step = map(self._reverse_index, (start, stop, step)) + idx = slice(start, stop, step) elif isinstance(idx, int): + if self.__reversed: + idx = self._reverse_index(idx) start = stop = idx else: raise TypeError('indices must be integers or slices') if start < 0 or stop < 0: # We need to consume the entire iterable to be able to slice from the end # Obviously, never use this with infinite iterables - self.exhaust() - else: - n = max(start, stop) - len(self.__cache) + 1 - if n > 0: - self.__cache.extend(itertools.islice(self.__iterable, n)) + return self.exhaust()[idx] + + n = max(start, stop) - len(self.__cache) + 1 + if n > 0: + self.__cache.extend(itertools.islice(self.__iterable, n)) return self.__cache[idx] def __bool__(self): try: - self[0] + self[-1] if self.__reversed else self[0] except IndexError: return False return True @@ -3996,6 +4010,17 @@ class LazyList(collections.Sequence): self.exhaust() return len(self.__cache) + def __reversed__(self): + self.__reversed = not self.__reversed + return self + + def __repr__(self): + # repr and str should mimic a list. So we exhaust the iterable + return repr(self.exhaust()) + + def __str__(self): + return repr(self.exhaust()) + class PagedList(object): def __len__(self): |