Tests and refinement

Author: mschauer

src/flatten.jl

@@ -1,6 +1,6 @@
 # Special case flatten of iterators of static arrays.
 import Base.Iterators: flatten_iteratorsize, flatten_length
-flatten_iteratorsize(::Union{Base.HasShape, Base.HasLength}, ::Type{<:Union{SArray,MArray}}) = Base.HasLength()
-function flatten_length(f, T::Type{<:Union{SArray,MArray}})
+flatten_iteratorsize(::Union{Base.HasShape, Base.HasLength}, ::Type{<:StaticArray}) = Base.HasLength()
+function flatten_length(f, T::Type{<:StaticArray})
     length(T)*length(f.it)
 end

test/flatten.jl

@@ -0,0 +1,13 @@
+using StaticArrays, Test
+
+@testset "Iterators.flatten" begin
+    for x in [SVector(1.0, 2.0), MVector(1.0, 2.0),
+            @SMatrix([1.0 2.0; 3.0 4.0]), @MMatrix([1.0 2.0]),
+            Size(1,2)([1.0 2.0])
+            ]
+        X = [x,x,x]
+        @test length(Iterators.flatten(X)) == length(X)*length(x)
+        @test collect(Iterators.flatten(typeof(x)[])) == []
+        @test collect(Iterators.flatten(X)) == [x..., x..., x...]
+    end
+end

test/runtests.jl

@@ -40,6 +40,7 @@ include("lu.jl")
 Random.seed!(42); include("qr.jl")
 Random.seed!(42); include("chol.jl") # hermitian_type(::Type{Any}) for block algorithm
 include("deque.jl")
+include("flatten.jl")
 include("io.jl")
 include("svd.jl")
 Random.seed!(42); include("fixed_size_arrays.jl")