iteratorsize for partition iterator

Author: mschauer

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)

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))
+    @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]