[utils] Improve LazyList

* Add `repr` and `str` that mimics `list`
* Add `reversed`. Unlike `[::-1]`, reversed does not exhaust the iterable and modifies the `LazyList` in-place
* Add tests
This commit is contained in:
pukkandan 2021-06-12 20:44:30 +05:30
parent 8ba8714880
commit 28419ca2c8
No known key found for this signature in database
GPG Key ID: 0F00D95A001F4698
2 changed files with 77 additions and 9 deletions

View File

@ -12,6 +12,7 @@
# Various small unit tests # Various small unit tests
import io import io
import itertools
import json import json
import xml.etree.ElementTree import xml.etree.ElementTree
@ -108,6 +109,7 @@
cli_bool_option, cli_bool_option,
parse_codecs, parse_codecs,
iri_to_uri, iri_to_uri,
LazyList,
) )
from yt_dlp.compat import ( from yt_dlp.compat import (
compat_chr, compat_chr,
@ -1525,6 +1527,47 @@ def test_clean_podcast_url(self):
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://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') 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__': if __name__ == '__main__':
unittest.main() unittest.main()

View File

@ -3954,10 +3954,14 @@ class LazyList(collections.Sequence):
def __init__(self, iterable): def __init__(self, iterable):
self.__iterable = iter(iterable) self.__iterable = iter(iterable)
self.__cache = [] self.__cache = []
self.__reversed = False
def __iter__(self): def __iter__(self):
for item in self.__cache: if self.__reversed:
yield item # 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: for item in self.__iterable:
self.__cache.append(item) self.__cache.append(item)
yield item yield item
@ -3965,21 +3969,31 @@ def __iter__(self):
def exhaust(self): def exhaust(self):
''' Evaluate the entire iterable ''' ''' Evaluate the entire iterable '''
self.__cache.extend(self.__iterable) self.__cache.extend(self.__iterable)
return self.__cache
@staticmethod
def _reverse_index(x):
return -(x + 1)
def __getitem__(self, idx): def __getitem__(self, idx):
if isinstance(idx, slice): if isinstance(idx, slice):
step = idx.step or 1 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 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): elif isinstance(idx, int):
if self.__reversed:
idx = self._reverse_index(idx)
start = stop = idx start = stop = idx
else: else:
raise TypeError('indices must be integers or slices') raise TypeError('indices must be integers or slices')
if start < 0 or stop < 0: if start < 0 or stop < 0:
# We need to consume the entire iterable to be able to slice from the end # We need to consume the entire iterable to be able to slice from the end
# Obviously, never use this with infinite iterables # Obviously, never use this with infinite iterables
self.exhaust() return self.exhaust()[idx]
else:
n = max(start, stop) - len(self.__cache) + 1 n = max(start, stop) - len(self.__cache) + 1
if n > 0: if n > 0:
self.__cache.extend(itertools.islice(self.__iterable, n)) self.__cache.extend(itertools.islice(self.__iterable, n))
@ -3987,7 +4001,7 @@ def __getitem__(self, idx):
def __bool__(self): def __bool__(self):
try: try:
self[0] self[-1] if self.__reversed else self[0]
except IndexError: except IndexError:
return False return False
return True return True
@ -3996,6 +4010,17 @@ def __len__(self):
self.exhaust() self.exhaust()
return len(self.__cache) 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): class PagedList(object):
def __len__(self): def __len__(self):