diff --git a/funcy/seqs.py b/funcy/seqs.py index 760a18e..d2e80fa 100644 --- a/funcy/seqs.py +++ b/funcy/seqs.py @@ -363,8 +363,12 @@ def _cut_iter(drop_tail, n, step, seq): if len(pool) < n: break yield pool - pool = pool[step:] - pool.extend(islice(it, step)) + if step < n: + pool = pool[step:] + pool.extend(islice(it, step)) + else: + take(step - n, it) # discard the gap between parts + pool = take(n, it) if not drop_tail: for item in _cut_seq(drop_tail, n, step, pool): yield item diff --git a/tests/test_seqs.py b/tests/test_seqs.py index 4a4c0e6..d0c8534 100644 --- a/tests/test_seqs.py +++ b/tests/test_seqs.py @@ -170,6 +170,18 @@ def test_chunks(): assert lchunks(2, 1, [0, 1, 2, 3]) == [[0, 1], [1, 2], [2, 3], [3]] assert lchunks(3, 1, iter(range(3))) == [[0, 1, 2], [1, 2], [2]] +def test_partition_step_gt_n(): + # The iterator path must skip step - n items between parts, like the + # Sequence path, when step > n. + assert lpartition(2, 3, iter(range(10))) == lpartition(2, 3, list(range(10))) + assert lpartition(2, 3, iter(range(10))) == [[0, 1], [3, 4], [6, 7]] + +def test_chunks_step_gt_n(): + # The iterator path must skip step - n items between parts, like the + # Sequence path, when step > n. + assert lchunks(2, 3, iter(range(10))) == lchunks(2, 3, list(range(10))) + assert lchunks(2, 3, iter(range(10))) == [[0, 1], [3, 4], [6, 7], [9]] + def test_partition_by(): assert lpartition_by(lambda x: x == 3, [1,2,3,4,5]) == [[1,2], [3], [4,5]] assert lpartition_by('x', 'abxcd') == [['a', 'b'], ['x'], ['c', 'd']]