Iteratorsize for partition iterator (#16552)

JuliaLang · Jul 12, 2017 · 0d83de9 · 0d83de9
1 parent 8678b5e
commit 0d83de9
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/base/iterators.jl b/base/iterators.jl
@@ -799,6 +799,11 @@ mutable struct PartitionIterator{T}
 end
 
 eltype(::Type{PartitionIterator{T}}) where {T} = Vector{eltype(T)}
+partition_iteratorsize(::HasShape) = HasLength()
+partition_iteratorsize(isz) = isz
+function iteratorsize(::Type{PartitionIterator{T}}) where {T}
+    partition_iteratorsize(iteratorsize(T))
+end
 
 function length(itr::PartitionIterator)
     l = length(itr.c)

diff --git a/test/iterators.jl b/test/iterators.jl
@@ -356,6 +356,13 @@ let v = collect(partition([1,2,3,4,5], 1))
     @test all(i->v[i][1] == i, v)
 end
 
+let v1 = collect(partition([1,2,3,4,5], 2)),
+    v2 = collect(partition(flatten([[1,2],[3,4],5]), 2)) # collecting partition with SizeUnkown
+    @test v1[1] == v2[1] == [1,2]
+    @test v1[2] == v2[2] == [3,4]
+    @test v1[3] == v2[3] == [5]
+end
+
 let v = collect(partition([1,2,3,4,5], 2))
     @test v[1] == [1,2]
     @test v[2] == [3,4]