diff --git a/base/array.jl b/base/array.jl index 15c354dce6085..7febf2a63828f 100644 --- a/base/array.jl +++ b/base/array.jl @@ -758,8 +758,10 @@ else end end -_array_for(::Type{T}, itr, ::HasLength) where {T} = Vector{T}(undef, Int(length(itr)::Integer)) -_array_for(::Type{T}, itr, ::HasShape{N}) where {T,N} = similar(Array{T,N}, axes(itr)) +_array_for(::Type{T}, itr, isz::HasLength) where {T} = _array_for(T, itr, isz, length(itr)) +_array_for(::Type{T}, itr, isz::HasShape{N}) where {T,N} = _array_for(T, itr, isz, axes(itr)) +_array_for(::Type{T}, itr, ::HasLength, len) where {T} = Vector{T}(undef, len) +_array_for(::Type{T}, itr, ::HasShape{N}, axs) where {T,N} = similar(Array{T,N}, axs) function collect(itr::Generator) isz = IteratorSize(itr.iter) @@ -767,12 +769,14 @@ function collect(itr::Generator) if isa(isz, SizeUnknown) return grow_to!(Vector{et}(), itr) else + shape = isz isa HasLength ? length(itr) : axes(itr) y = iterate(itr) if y === nothing return _array_for(et, itr.iter, isz) end v1, st = y - collect_to_with_first!(_array_for(typeof(v1), itr.iter, isz), v1, itr, st) + arr = _array_for(typeof(v1), itr.iter, isz, shape) + return collect_to_with_first!(arr, v1, itr, st) end end diff --git a/test/iterators.jl b/test/iterators.jl index fb8edcab92209..c7d00c4e7e2e8 100644 --- a/test/iterators.jl +++ b/test/iterators.jl @@ -292,6 +292,15 @@ let (a, b) = (1:3, [4 6; end end +# collect stateful iterator +let + itr = (i+1 for i in Base.Stateful([1,2,3])) + @test collect(itr) == [2, 3, 4] + A = zeros(Int, 0, 0) + itr = (i-1 for i in Base.Stateful(A)) + @test collect(itr) == Int[] # Stateful do not preserve shape +end + # with 1D inputs let a = 1:2, b = 1.0:10.0, @@ -860,4 +869,4 @@ end @test Iterators.peel(1:10)[2] |> collect == 2:10 @test Iterators.peel(x^2 for x in 2:4)[1] == 4 @test Iterators.peel(x^2 for x in 2:4)[2] |> collect == [9, 16] -end \ No newline at end of file +end