diff --git a/base/iterators.jl b/base/iterators.jl index b329a6c6f9eff..6c14cfe310c06 100644 --- a/base/iterators.jl +++ b/base/iterators.jl @@ -264,6 +264,33 @@ struct Zip{Is<:Tuple} is::Is end +""" + truncate(s1, s2) + +Check two array shapes for compatibility, and return whichever shape has less dimensions. + +# Examples +```jldoctest +julia> Iterators.truncate([1,2,3], 1:10) +(Base.OneTo(3),) + +julia> Iterators.truncate([2,3,1], [2,3,5,6,4]) +(Base.OneTo(3),) +``` +""" + +function truncate(a::AbstractArray, b::AbstractArray) + truncate(axes(a), axes(b)) +end + +function truncate(a::Tuple{Base.OneTo{Int}}, b::Tuple{Base.OneTo{Int}}) + if length(a[1]) < length(b[1]) + return a + else + return b + end +end + """ zip(iters...) @@ -313,7 +340,9 @@ function _zip_min_length(is) end _zip_min_length(is::Tuple{}) = nothing size(z::Zip) = _promote_shape(map(size, z.is)...) -axes(z::Zip) = _promote_shape(map(axes, z.is)...) +axes(z::Zip) = truncate(map(axes, z.is)...) +truncate(p, q...) = truncate(p, truncate(q...)) +truncate(p) = p _promote_shape(a, b...) = promote_shape(a, _promote_shape(b...)) _promote_shape(a) = a eltype(::Type{Zip{Is}}) where {Is<:Tuple} = _zip_eltype(Is) diff --git a/test/iterators.jl b/test/iterators.jl index 86b2cd1d2969d..1ffd4fb5ff1ab 100644 --- a/test/iterators.jl +++ b/test/iterators.jl @@ -76,6 +76,13 @@ end @test_throws MethodError collect(rest(countfrom(1), 5)) +# truncate +# -------- +@test Iterators.truncate([1,2,3], [2,3,4,5]) == (Base.OneTo(3),) +@test Iterators.truncate(1:10, 1:50) == (Base.OneTo(10),) +@test Iterators.truncate([1,2,3], 1:10) == (Base.OneTo(3),) +@test Iterators.truncate((Base.OneTo(3),), (Base.OneTo(5),)) == (Base.OneTo(3),) + # countfrom # --------- let i = 0, k = 1