diff --git a/base/iterator.jl b/base/iterator.jl index 833222ab02e19f..cf0170a5bb84ff 100644 --- a/base/iterator.jl +++ b/base/iterator.jl @@ -341,3 +341,49 @@ function collect{I<:IteratorND}(g::Generator{I}) dest[1] = first return map_to!(g.f, 2, st, dest, g.iter) end + + +""" + collect(n, collection) + +Returns an `n` elements at a time iterator over collection + +```jldoctest +julia> collect(collect(2, [1,2,3,4,5])) +3-element Array{Array{Int64,1},1}: + [1,2] + [3,4] + [5] +``` +""" +function collect{T}(n::Int, c::T) + @assert n > 0 + return CollectIterator{T}(n, c) +end + +type CollectIterator{T} + n::Int + c::T +end + +eltype{T}(::Type{CollectIterator{T}}) = Vector{eltype(T)} + +start(itr::CollectIterator) = start(itr.c) + +done(itr::CollectIterator, state) = done(itr.c, state) + +function next{T<:Vector}(itr::CollectIterator{T}, state) + l = state + r = min(state + itr.n-1, length(itr.c)) + return itr.c[l:r], r + 1 +end + +function next(itr::CollectIterator, state) + v = Vector{eltype(itr.c)}(itr.n) + i = 0 + while !done(itr.c, state) && i < itr.n + i += 1 + v[i], state = next(itr.c, state) + end + return resize!(v, i), state +end diff --git a/doc/stdlib/collections.rst b/doc/stdlib/collections.rst index 5edb9ef2bae792..0571f9185d944d 100644 --- a/doc/stdlib/collections.rst +++ b/doc/stdlib/collections.rst @@ -639,6 +639,8 @@ Iterable Collections Return an array of type ``Array{element_type,1}`` of all items in a collection. +.. function:: collect(n, collection) + .. function:: issubset(a, b) ⊆(a,b) -> Bool ⊈(a,b) -> Bool diff --git a/test/functional.jl b/test/functional.jl index 826e4e8d2792c3..a51226eb43711f 100644 --- a/test/functional.jl +++ b/test/functional.jl @@ -204,3 +204,31 @@ let f(g) = (@test size(g.iter)==(2,3)) end @test_throws DimensionMismatch Base.IteratorND(1:2, (2,3)) + + +# collect(n, c) + +let v = collect(collect(1, [1,2,3,4,5])) + @test all(i->v[i][1] == i, v) +end + +let v = collect(collect(2, [1,2,3,4,5])) + @test v[1] == [1,2] + @test v[2] == [3,4] + @test v[3] == [5] +end + +let v = collect(collect(3, enumerate([1,2,3,4,5]))) + @test v[1] == [(1,1),(2,2),(3,3)] + @test v[2] == [(4,4),(5,5)] +end + +for n in [5,6] + @test collect(collect(n, [1,2,3,4,5]))[1] == [1,2,3,4,5] + @test collect(collect(n, enumerate([1,2,3,4,5])))[1] == + [(1,1),(2,2),(3,3),(4,4),(5,5)] +end + + +@test_throws AssertionError collect(0, [1,2,3,4,5]) +@test_throws AssertionError collect(-1, [1,2,3,4,5])