Permalink
Browse files

update partition for new protocol

  • Loading branch information...
iamed2 committed May 31, 2018
1 parent b522758 commit 66d1951c2634d06ab629fcf35ea227ddbe9a2342
Showing with 29 additions and 32 deletions.
  1. +25 −32 src/IterTools.jl
  2. +4 −0 test/runtests.jl
@@ -392,62 +392,55 @@ i = (4, 5)
i = (7, 8) i = (7, 8)
``` ```
""" """
function partition(xs::I, n::Int) where I @inline partition(xs, n::Int) = partition(xs, n, n)
Partition{I, n}(xs, n)
end
function partition(xs::I, n::Int, step::Int) where I function partition(xs::I, n::Int, step::Int) where I
if step < 1 if step < 1
throw(ArgumentError("Partition step must be at least 1.")) throw(ArgumentError("Partition step must be at least 1."))
end end
if n < 1
throw(ArgumentError("Partition size n must be at least 1"))
end
Partition{I, n}(xs, step) Partition{I, n}(xs, step)
end end
function start(it::Partition{I, N}) where {I, N} function iterate(it::Partition{I, N}, state=nothing) where {I, N}
p = Vector{eltype(I)}(undef, N) if state === nothing
s = start(it.xs) result = Vector{eltype(I)}(undef, N)
for i in 1:(N - 1)
if done(it.xs, s) result[1], xs_state = @something iterate(it.xs)
break
for i in 2:N
result[i], xs_state = @something iterate(it.xs, xs_state)
end end
(p[i], s) = next(it.xs, s) else
(xs_state, result) = state
result[end], xs_state = @something iterate(it.xs, xs_state)
end end
(s, p)
end
function next(it::Partition{I, N}, state) where {I, N} p = similar(result)
(s, p0) = state
(x, s) = next(it.xs, s)
ans = p0; ans[end] = x
p = similar(p0)
overlap = max(0, N - it.step) overlap = max(0, N - it.step)
for i in 1:overlap p[1:overlap] .= result[it.step .+ (1:overlap)]
p[i] = ans[it.step + i]
end
# when step > n, skip over some elements # when step > n, skip over some elements
for i in 1:max(0, it.step - N) for i in 1:max(0, it.step - N)
if done(it.xs, s) xs_iter = iterate(it.xs, xs_state)
break xs_iter === nothing && break
end _, xs_state = xs_iter
(x, s) = next(it.xs, s)
end end
for i in (overlap + 1):(N - 1) for i in (overlap + 1):(N - 1)
if done(it.xs, s) xs_iter = iterate(it.xs, xs_state)
break xs_iter === nothing && break
end
(x, s) = next(it.xs, s) p[i], xs_state = xs_iter
p[i] = x
end end
(tuple(ans...), (s, p)) return (tuple(result...)::eltype(Partition{I, N}), (xs_state, p))
end end
done(it::Partition, state) = done(it.xs, state[1])
# Group output from an iterator based on a key function. # Group output from an iterator based on a key function.
# Consecutive entries from the iterator with the same # Consecutive entries from the iterator with the same
@@ -143,6 +143,10 @@ include("testing_macros.jl")
@test eltype(pa3) == Tuple{Int, Int} @test eltype(pa3) == Tuple{Int, Int}
@test collect(pa3) == [] @test collect(pa3) == []
pa4 = partition(1:8, 1)
@test eltype(pa4) == Tuple{Int}
@test collect(pa4) == [(1,), (2,), (3,), (4,), (5,), (6,), (7,), (8,)]
@test_throws ArgumentError partition(take(countfrom(1), 8), 2, 0) @test_throws ArgumentError partition(take(countfrom(1), 8), 2, 0)
end end

0 comments on commit 66d1951

Please sign in to comment.