Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make partition support non-1-indexed vector. #40830

Merged
merged 15 commits into from
May 21, 2021
12 changes: 6 additions & 6 deletions base/iterators.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1171,15 +1171,15 @@ function length(itr::PartitionIterator)
return cld(l, itr.n)
end

function iterate(itr::PartitionIterator{<:AbstractRange}, state=1)
state > length(itr.c) && return nothing
r = min(state + itr.n - 1, length(itr.c))
function iterate(itr::PartitionIterator{<:AbstractRange}, state = firstindex(itr.c))
state > lastindex(itr.c) && return nothing
r = min(state + itr.n - 1, lastindex(itr.c))
return @inbounds itr.c[state:r], r + 1
end

function iterate(itr::PartitionIterator{<:AbstractArray}, state=1)
state > length(itr.c) && return nothing
r = min(state + itr.n - 1, length(itr.c))
function iterate(itr::PartitionIterator{<:AbstractArray}, state = firstindex(itr.c))
state > lastindex(itr.c) && return nothing
r = min(state + itr.n - 1, lastindex(itr.c))
return @inbounds view(itr.c, state:r), r + 1
end

Expand Down
5 changes: 5 additions & 0 deletions test/iterators.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

using Base.Iterators
using Random
using Base: IdentityUnitRange

@test Base.IteratorSize(Any) isa Base.SizeUnknown

Expand Down Expand Up @@ -848,3 +849,7 @@ end
@test cumprod(x + 1 for x in 1:3) == [2, 6, 24]
@test accumulate(+, (x^2 for x in 1:3); init=100) == [101, 105, 114]
end

@testset "proper patition for non-1-indexed vector" begin
@test partition(IdentityUnitRange(11:19), 5) |> collect == [11:15,16:19] # IdentityUnitRange
end
9 changes: 9 additions & 0 deletions test/offsetarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

isdefined(Main, :OffsetArrays) || @eval Main include("testhelpers/OffsetArrays.jl")
using .Main.OffsetArrays
import .Main.OffsetArrays: IdOffsetRange
using DelimitedFiles
using Random
using LinearAlgebra
Expand Down Expand Up @@ -786,3 +787,11 @@ end
@test b[i] == a[r[i]]
end
end

@testset "proper patition for non-1-indexed vector" begin
@test Iterators.partition(OffsetArray(1:10,10), 5) |> collect == [1:5,6:10] # OffsetVector
N5N3 marked this conversation as resolved.
Show resolved Hide resolved
@test Iterators.partition(OffsetArray(collect(1:10),10), 5) |> collect == [1:5,6:10] # OffsetVector
@test Iterators.partition(OffsetArray(reshape(1:9,3,3), (3,3)), 5) |> collect == [1:5,6:9] #OffsetMatrix
@test Iterators.partition(OffsetArray(reshape(collect(1:9),3,3), (3,3)), 5) |> collect == [1:5,6:9] #OffsetMatrix
@test Iterators.partition(IdOffsetRange(2:7,10), 5) |> collect == [12:16,17:17] # IdOffsetRange
end