Iterate sequentially over several collections

I have two collections with the same eltype, say c1 = 1:3 and c2 = 5:7. I would like to be able to write a generator over the joined collections, something like

(i for i in joincol(c1, c2))

so that this becomes a generator of [1,2,3,5,6,7] without allocating [c1..., c2...] itself

Is there some equivalent of the joincol function in Base, or some standard way to do it without writing a custom iterate method?

Iterators.flatten((1:3, 5:7))

1 Like

Fantastic, thanks!

Obviously I need to study up on the already implemented iterators in the Iterators module, but since I already took the last 5 minutes doing this, I’ll post it anyway to feel like it was worth it :stuck_out_tongue:

struct Chain{T}
    t::T # Tuple of iterators
end
Chain(x...) = Chain(x)
Base.IteratorSize(::Type{Chain{T}}) where {T} = all(x->Base.IteratorSize(x)==Base.HasLength() || typeof(Base.IteratorSize(x)) <: Base.HasShape, T.parameters) ? Base.HasLength() : Base.SizeUnknown()
Base.length(c::C) where {C <: Chain} = Base.IteratorSize(C) == Base.HasLength() ? sum(x->length(x), c.t) : missing

function Base.iterate(c::Chain)
    st = nothing
    itr = c.t[1]
    i = 1
    while st === nothing && i <= length(c.t)
        itr = c.t[i]
        st = iterate(itr)
        i += 1
    end
    st === nothing && return nothing
    val, s = st
    return val, (itr, s, i)
end

function Base.iterate(c::Chain, state)
    itr, s, i = state
    st = iterate(itr, s)
    while st === nothing && i <= length(c.t)
        itr = c.t[i]
        st = iterate(itr)
        i += 1
    end
    st === nothing && return nothing
    val, s = st
    return val, (itr, s, i)
end

It’s probably not quite perfect, but whatever, that’s what I came up w/ in 5 minutes.

Haha! This was the kind of thing I was going to attempt before I posted. Sorry for stealing 5 mins of your time for this. Github needs you. Go! :stuck_out_tongue_winking_eye:

I needed and found flatten a week or so ago. It took a lot of digging to find, more than usual.

Can I ask why you wouldn’t want to do the following?

(i for i in [c1;c2])

This allocates an array. flatten is lazy so it doesn’t allocate, just generates elements from the first generator then the second.

1 Like