diff --git a/src/qwlist/eager.py b/src/qwlist/eager.py index 46ab843..32f05cb 100644 --- a/src/qwlist/eager.py +++ b/src/qwlist/eager.py @@ -399,5 +399,24 @@ def inner(): yield elem return EagerQList(inner()) + def take_while(self, pred: Callable[[T], bool]) -> "Lazy[T]": + """ + Creates a new EagerQList that contains elements based on a predicate. Takes a function as an argument. + It will call this function on each element of the iterator, and yield elements while the function + returns `True`. After `False` is returned, iteration stops, and the rest of the elements is ignored. + + Args: + pred (Callable[[T], bool]): `function (T) -> bool` + + Returns: + `EagerQList[T]` + """ + def inner(): + for elem in self: + if not pred(elem): + return + yield elem + return EagerQList(inner()) + if __name__ == '__main__': EagerQList([[1, 2], 3]).flatten().foreach(print) \ No newline at end of file diff --git a/src/qwlist/qwlist.py b/src/qwlist/qwlist.py index 63a91c4..367982a 100644 --- a/src/qwlist/qwlist.py +++ b/src/qwlist/qwlist.py @@ -492,8 +492,9 @@ def sum(self, init: Optional[T] = None) -> T: def take_while(self, pred: Callable[[T], bool]) -> "Lazy[T]": """ Creates a new Lazy that yields elements based on a predicate. Takes a function as an argument. + Creates a new Lazy that yields elements based on a predicate. Takes a function as an argument. It will call this function on each element of the iterator, and yield elements while the function - returns `True`. After `False` is returned, iteration stops, and the rest of the elements are ignored. + returns `True`. After `False` is returned, iteration stops, and the rest of the elements is ignored. Args: pred (Callable[[T], bool]): `function (T) -> bool` @@ -1006,6 +1007,25 @@ def inner(): yield elem return Lazy(inner()) + def take_while(self, pred: Callable[[T], bool]) -> "Lazy[T]": + """ + Creates a Lazy that yields elements based on a predicate. Takes a function as an argument. + It will call this function on each element of the iterator, and yield elements while the function + returns `True`. After `False` is returned, iteration stops, and the rest of the elements is ignored. + + Args: + pred (Callable[[T], bool]): `function (T) -> bool` + + Returns: + `Lazy[T]` + """ + def inner(): + for elem in self: + if not pred(elem): + return + yield elem + return Lazy(inner()) + if __name__ == '__main__': def naturals(start): @@ -1014,10 +1034,14 @@ def naturals(start): yield current current += 1 - for n in Lazy(naturals(0)).take(10): - print(n) - - primes = Lazy(naturals(2)).filter(lambda n: Lazy(naturals(2)).take_while(lambda p: p * p <= n).all(lambda x: n % x != 0)) - for p in primes.take(100).skip(80): + primes = ( + Lazy(naturals(2)) + .filter(lambda n: ( + Lazy(naturals(2)) + .take_while(lambda p: p * p <= n) + .all(lambda x: n % x != 0) + )) + ) + for p in primes.skip(1000).take(10): print(p) diff --git a/src/todo.md b/src/todo.md index dd4cc6d..ce7c939 100644 --- a/src/todo.md +++ b/src/todo.md @@ -1,8 +1,8 @@ QList, Lazy, EagerQList ### take_while -[ ] [ ] [ ] implemented \ -[ ] [ ] [ ] documented \ -[ ] [ ] [ ] tested +[x] [x] [x] implemented \ +[x] [x] [x] documented \ +[x] [x] [x] tested ### all [x] [x] [x] implemented \ diff --git a/tests/test_eager.py b/tests/test_eager.py index 4c02623..cef071e 100644 --- a/tests/test_eager.py +++ b/tests/test_eager.py @@ -305,6 +305,7 @@ def test_all(): assert not EagerQList(['', 'a', 'aa']).all() assert EagerQList(range(10)).filter(lambda x: x % 2 == 1).map(lambda x: x * 2).all(lambda x: x % 2 == 0) + def test_any(): assert EagerQList([1, True, [1, 2, 3]]).any() assert not EagerQList().any() @@ -313,4 +314,22 @@ def test_any(): assert EagerQList(['abc', 'def', 'gdi']).any(mapper=lambda s: len(s) > 1) assert not EagerQList([False, False, False]).any() assert EagerQList(['', 'a', 'aa']).any() - assert EagerQList(range(10)).filter(lambda x: x < 5).all(lambda x: x % 2 == 0) + assert EagerQList(range(10)).filter(lambda x: x < 5).any(lambda x: x % 2 == 0) + + +def test_take_while(): + expected = EagerQList([0, 1, 2]) + res = EagerQList(range(10)).take_while(lambda n: n < 3) + assert res == expected + + expected = EagerQList() + res = EagerQList(range(10)).take_while(lambda n: isinstance(n, str)) + assert res == expected + + expected = EagerQList() + res = EagerQList([]).take_while(lambda x: x > 2) + assert res == expected + + expected = EagerQList(range(10)) + res = EagerQList(range(5)).take_while(lambda x: x < 100).chain([5, 6, 7, 8, 9]) + assert res == expected diff --git a/tests/test_lazy.py b/tests/test_lazy.py index 1a45c5a..da17e61 100644 --- a/tests/test_lazy.py +++ b/tests/test_lazy.py @@ -374,6 +374,7 @@ def test_all(): assert not Lazy([False, False, False]).all() assert not Lazy(['', 'a', 'aa']).all() + def test_any(): assert Lazy([1, True, [1, 2, 3]]).any() assert not Lazy([]).any() @@ -381,4 +382,22 @@ def test_any(): assert Lazy([True, True, False]).any(mapper=lambda x: not x) assert Lazy(['abc', 'def', 'gdi']).any(mapper=lambda s: len(s) > 1) assert not Lazy([False, False, False]).any() - assert Lazy(['', 'a', 'aa']).any() \ No newline at end of file + assert Lazy(['', 'a', 'aa']).any() + + +def test_take_while(): + expected = QList([0, 1, 2]) + res = Lazy(range(10)).take_while(lambda n: n < 3).collect() + assert res == expected + + expected = QList() + res = Lazy(range(10)).take_while(lambda n: isinstance(n, str)).collect() + assert res == expected + + expected = QList() + res = Lazy([]).take_while(lambda x: x > 2).collect() + assert res == expected + + expected = QList(range(10)) + res = Lazy(range(5)).take_while(lambda x: x < 100).chain([5, 6, 7, 8, 9]).collect() + assert res == expected \ No newline at end of file diff --git a/tests/test_qwlist.py b/tests/test_qwlist.py index 09ac8a8..bd785d7 100644 --- a/tests/test_qwlist.py +++ b/tests/test_qwlist.py @@ -289,7 +289,7 @@ def test_flatten(): assert expected == res try: - QList([1, 2, 3]).flatten() + QList([1, 2, 3]).flatten().collect() except TypeError: assert True else: @@ -413,6 +413,7 @@ def test_all(): assert not QList(['', 'a', 'aa']).all() assert QList(range(10)).filter(lambda x: x % 2 == 1).map(lambda x: x * 2).all(lambda x: x % 2 == 0) + def test_any(): assert QList([1, True, [1, 2, 3]]).any() assert not QList().any() @@ -421,5 +422,23 @@ def test_any(): assert QList(['abc', 'def', 'gdi']).any(mapper=lambda s: len(s) > 1) assert not QList([False, False, False]).any() assert QList(['', 'a', 'aa']).any() - assert QList(range(10)).filter(lambda x: x < 5).all(lambda x: x % 2 == 0) + assert QList(range(10)).filter(lambda x: x < 5).any(lambda x: x % 2 == 0) + + +def test_take_while(): + expected = QList([0, 1, 2]) + res = QList(range(10)).take_while(lambda n: n < 3).collect() + assert res == expected + + expected = QList() + res = QList(range(10)).take_while(lambda n: isinstance(n, str)).collect() + assert res == expected + + expected = QList() + res = QList().take_while(lambda x: x > 2).collect() + assert res == expected + + expected = QList(range(10)) + res = QList(range(5)).take_while(lambda x: x < 100).chain([5, 6, 7, 8, 9]).collect() + assert res == expected