add chunk()

Author: cleoold

RunTest.ps1

@@ -1,4 +1,4 @@
-coverage run -m pytest
+coverage run -m pytest @args
 if ($?)
 {
     coverage report -m

doc/api/types_linq.enumerable.rst

@@ -506,6 +506,35 @@ Example
 
 ----
 
+instancemethod ``chunk(size)``
+--------------------------------
+
+Parameters
+  - `size` (``int``)
+
+Returns
+  - ``Enumerable[MutableSequence[TSource_co]]``
+
+Splits the elements of a sequence into chunks of size at most the provided size. Raises
+`InvalidOperationError` if `size` is less than 1.
+
+Example
+    .. code-block:: python
+
+        >>> def source(i):
+        ...     while True:
+        ...         yield i
+        ...         i *= 3
+
+        >>> en = Enumerable(source(1)).chunk(4).take(3)
+        >>> for chunk in en:
+        ...     print(chunk)
+        [1, 3, 9, 27]
+        [81, 243, 729, 2187]
+        [6561, 19683, 59049, 177147]
+
+----
+
 instancemethod ``concat(second)``
 -----------------------------------
 

doc/api_spec.py

@@ -52,6 +52,7 @@ class ClassSpec(TypedDict):
                     'average',
                     'average2',
                     'cast',
+                    'chunk',
                     'concat',
                     'contains',
                     'count',

tests/test_usage.py

@@ -212,6 +212,48 @@ def test_average2_overload2(self):
         assert avg == 5.8
 
 
+class TestChunkMethod:
+    def test_infinite_exact(self):
+        def source(i):
+            while True:
+                yield i
+                i *= 3
+        en = Enumerable(source(1)).chunk(4).take(3)
+        assert en.to_list() == [
+            [1, 3, 9, 27], [81, 243, 729, 2187], [6561, 19683, 59049, 177147],
+        ]
+
+    def test_sized_last_not_full(self):
+        ints = [1, 3, 9, 27, 81, 243, 729]
+        en = Enumerable(ints).chunk(3)
+        assert en.to_list() == [[1, 3, 9], [27, 81, 243], [729]]
+
+    def test_elements_less_than_size(self):
+        ints = [1, 3, 9, 27]
+        en = Enumerable(ints).chunk(5)
+        assert en.to_list() == [[1, 3, 9, 27]]
+
+    def test_empty(self):
+        en = Enumerable([])
+        assert en.chunk(3).to_list() == []
+
+    def test_mutate_source_before_iteration(self):
+        ints = [1, 3, 9, 27, 81, 243, 729]
+        en = Enumerable(ints).chunk(2)
+        it = iter(en)
+        assert next(it) == [1, 3]
+        ints.remove(27)
+        assert next(it) == [9, 81]
+        assert next(it) == [243, 729]
+
+    def test_invalid_size(self):
+        en = Enumerable([0])
+        with pytest.raises(InvalidOperationError):
+            en.chunk(0)
+        with pytest.raises(InvalidOperationError):
+            en.chunk(-1)
+
+
 class TestConcatMethod:
     def test_concat(self):
         en1 = Enumerable([1, 2, 3])
@@ -371,7 +413,8 @@ def test_fallback(self):
             en3[1]
 
     def test_not_fallback_and_with_default(self):
-        en = Enumerable([])
+        en = Enumerable([1])
+        assert en[0, 'haha'] == 1
         assert en[1, 'haha'] == 'haha'
 
 

types_linq/enumerable.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Any, Callable, Container, Deque, Dict, Iterable, Iterator, List, NoReturn, Optional, Reversible, Sequence, Set, Sized, TYPE_CHECKING, Tuple, Type, Generic, Union
+from typing import Any, Callable, Container, Deque, Dict, Iterable, Iterator, List, MutableSequence, NoReturn, Optional, Reversible, Sequence, Set, Sized, TYPE_CHECKING, Tuple, Type, Generic, Union
 
 if TYPE_CHECKING:
     from .lookup import Lookup
@@ -255,6 +255,25 @@ def average2(self, *args):
     def cast(self, _: Type[TResult]) -> Enumerable[TResult]:
         return self  # type: ignore
 
+    def chunk(self, size: int) -> Enumerable[MutableSequence[TSource_co]]:
+        if size < 1:
+            raise InvalidOperationError('size must be greater than 0')
+        def inner():
+            lst: List[Any] = [1] * size
+            i = 0
+            for elem in self:
+                lst[i] = elem
+                if i == size - 1:
+                    yield lst
+                    i = 0
+                    lst = [1] * size
+                else:
+                    i += 1
+            if i > 0:
+                del lst[i:]
+                yield lst
+        return Enumerable(inner)
+
     def concat(self, second: Iterable[TSource_co]) -> Enumerable[TSource_co]:
         def inner():
             yield from self

types_linq/enumerable.pyi

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, Generic, Iterable, Iterator, List, NoReturn, Optional, Sequence, Set, Tuple, Type, Union, overload
+from typing import Any, Callable, Dict, Generic, Iterable, Iterator, List, MutableSequence, NoReturn, Optional, Sequence, Set, Tuple, Type, Union, overload
 
 from .lookup import Lookup
 from .grouping import Grouping
@@ -389,6 +389,28 @@ class Enumerable(Sequence[TSource_co], Generic[TSource_co]):
                 same_query: Enumerable[int] = query.cast(int)
         '''
 
+    # nonstanard: .NET returns fix-size arrays as elements instead
+    def chunk(self, size: int) -> Enumerable[MutableSequence[TSource_co]]:
+        '''
+        Splits the elements of a sequence into chunks of size at most the provided size. Raises
+        `InvalidOperationError` if `size` is less than 1.
+
+        Example
+            .. code-block:: python
+
+                >>> def source(i):
+                ...     while True:
+                ...         yield i
+                ...         i *= 3
+
+                >>> en = Enumerable(source(1)).chunk(4).take(3)
+                >>> for chunk in en:
+                ...     print(chunk)
+                [1, 3, 9, 27]
+                [81, 243, 729, 2187]
+                [6561, 19683, 59049, 177147]
+        '''
+
     def concat(self, second: Iterable[TSource_co]) -> Enumerable[TSource_co]:
         '''
         Concatenates two sequences.