Skip to content

Commit

Permalink
Add collect(n, collection)
Browse files Browse the repository at this point in the history
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]
```
  • Loading branch information
samoconnor committed Mar 10, 2016
1 parent 821679e commit 0483198
Show file tree
Hide file tree
Showing 3 changed files with 76 additions and 0 deletions.
46 changes: 46 additions & 0 deletions base/iterator.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
2 changes: 2 additions & 0 deletions doc/stdlib/collections.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
28 changes: 28 additions & 0 deletions test/functional.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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])

0 comments on commit 0483198

Please sign in to comment.