add take_while()

Author: cleoold

tests/test_usage.py

@@ -1259,6 +1259,38 @@ def test_count_zero_or_negative(self):
         assert en.take_last(-1).to_list() == []
 
 
+class TestTakeWhileMethod:
+    def test_take_while_some(self):
+        lst = ['1', '3', '5', '7', '', '1', '4', '5']
+        en = Enumerable(lst)
+        q = en.take_while(lambda x: x != '')
+        assert q.to_list() == ['1', '3', '5', '7']
+
+    def test_take_while_all(self):
+        lst = ['1', '3', '5', '7', '', '1', '4', '5']
+        en = Enumerable(lst)
+        q = en.take_while(lambda x: x != '77')
+        assert q.to_list() == lst
+
+    def test_take_while_nothing(self):
+        lst = ['1', '3', '5', '7', '', '1', '4', '5']
+        en = Enumerable(lst)
+        q = en.take_while(lambda x: not isinstance(x, str))
+        assert q.to_list() == []
+
+    def test_take_while2_some(self):
+        lst = ['1', '3', '5', '7', '', '1', '4', '5']
+        en = Enumerable(lst)
+        q = en.take_while2(lambda x, i: isinstance(x, str) and i < 4)
+        assert q.to_list() == ['1', '3', '5', '7']
+
+    def test_take_while2_all(self):
+        lst = ['1', '3', '5', '7', '', '1', '4', '5']
+        en = Enumerable(lst)
+        q = en.take_while2(lambda x, i: isinstance(x, str) and i < 8)
+        assert q.to_list() == lst
+
+
 class TestUnionMethod:
     def test_union(self):
         gen = (i for i in range(5))

types_linq/enumerable.py

@@ -778,6 +778,22 @@ def inner():
             yield from q
         return Enumerable(inner)
 
+    def take_while(self, predicate: Callable[[TSource_co], bool]) -> Enumerable[TSource_co]:
+        def inner():
+            for elem in self:
+                if not predicate(elem):
+                    break
+                yield elem
+        return Enumerable(inner)
+
+    def take_while2(self, predicate: Callable[[TSource_co, int], bool]) -> Enumerable[TSource_co]:
+        def inner():
+            for i, elem in enumerate(self):
+                if not predicate(elem, i):
+                    break
+                yield elem
+        return Enumerable(inner)
+
     def to_dict(self,
         key_selector: Callable[[TSource_co], TKey],
         *args: Callable[[TSource_co], TValue],

types_linq/enumerable.pyi

@@ -645,7 +645,16 @@ class Enumerable(Sequence[TSource_co], Generic[TSource_co]):
         Returns a new sequence that contains the last `count` elements.
         '''
 
-    # @@@ TODO
+    def take_while(self, predicate: Callable[[TSource_co], bool]) -> Enumerable[TSource_co]:
+        '''
+        Returns elements from the sequence as long as the condition is true and skips the remaining.
+        '''
+
+    def take_while2(self, predicate: Callable[[TSource_co, int], bool]) -> Enumerable[TSource_co]:
+        '''
+        Returns elements from the sequence as long as the condition is true and skips the remaining. The
+        element's index is used in the predicate function.
+        '''
 
     @overload
     def to_dict(self,