diff --git a/base/abstractarraymath.jl b/base/abstractarraymath.jl index e5ca36f71ac0a..5b04fa310e740 100644 --- a/base/abstractarraymath.jl +++ b/base/abstractarraymath.jl @@ -407,3 +407,40 @@ _reperr(s, n, N) = throw(ArgumentError("number of " * s * " repetitions " * return R end + +""" + eachrow(A::AbstractVecOrMat) + +Create a generator that iterates over the first dimension of vector or matrix `A`, +returning the rows as views. +See also [`eachcol`](@ref) and [`eachslice`](@ref). +""" +eachrow(A::AbstractVecOrMat) = (view(A, i, :) for i in axes(A, 1)) + + +""" + eachcol(A::AbstractVecOrMat) + +Create a generator that iterates over the second dimension of matrix `A`, returning the +columns as views. +See also [`eachrow`](@ref) and [`eachslice`](@ref). +""" +eachcol(A::AbstractVecOrMat) = (view(A, :, i) for i in axes(A, 2)) + +""" + eachslice(A::AbstractArray; dims) + +Create a generator that iterates over dimensions `dims` of `A`, returning views that select all +the data from the other dimensions in `A`. + +Only a single dimension in `dims` is currently supported. Equivalent to `(view(A,:,:,...,i,:,: +...)) for i in axes(A, dims))`, where `i` is in position `dims`. +See also [`eachrow`](@ref), [`eachcol`](@ref), and [`selectdim`](@ref). +""" +@inline function eachslice(A::AbstractArray; dims) + length(dims) == 1 || throw(ArgumentError("only single dimensions are supported")) + dim = first(dims) + dim <= ndims(A) || throw(DimensionMismatch("A doesn't have $dim dimensions")) + idx1, idx2 = ntuple(d->(:), dim-1), ntuple(d->(:), ndims(A)-dim) + return (view(A, idx1..., i, idx2...) for i in axes(A, dim)) +end diff --git a/base/exports.jl b/base/exports.jl index 784ac8b3ff207..d9c3812b5aebf 100644 --- a/base/exports.jl +++ b/base/exports.jl @@ -365,7 +365,10 @@ export cumsum!, accumulate, accumulate!, + eachcol, eachindex, + eachrow, + eachslice, extrema, fill!, fill, diff --git a/test/arrayops.jl b/test/arrayops.jl index a5e4b67715431..7790cabb036e9 100644 --- a/test/arrayops.jl +++ b/test/arrayops.jl @@ -1971,6 +1971,22 @@ end @test IndexStyle(selectdim(A, 3, 1)) == IndexStyle(view(A, :, :, 1)) == IndexLinear() end +# row/column/slice iterator tests +using Base: eachrow, eachcol +@testset "row/column/slice iterators" begin + # Simple ones + M = [1 2 3; 4 5 6; 7 8 9] + @test collect(eachrow(M)) == collect(eachslice(M, dims = 1)) == [[1, 2, 3], [4, 5, 6], [7, 8, 9]] + @test collect(eachcol(M)) == collect(eachslice(M, dims = 2)) == [[1, 4, 7], [2, 5, 8], [3, 6, 9]] + @test_throws DimensionMismatch eachslice(M, dims = 4) + + # Higher-dimensional case + M = reshape([(1:16)...], 2, 2, 2, 2) + @test_throws MethodError collect(eachrow(M)) + @test_throws MethodError collect(eachcol(M)) + @test collect(eachslice(M, dims = 1))[1][:, :, 1] == [1 5; 3 7] +end + ### ### IndexCartesian workout ###