A stream-like wrapper around iterable Julia objects. Based on https://github.com/andrewcooke/StatefulIterators.jl

StatefulIterators.jl

module StatefulIterators_

export StatefulIterator, available, peek, reset!
import Base: eltype, start, next, done, read, copy

type StatefulIterator{T, S}
    iter::T
    state::S
end
StatefulIterator(itr) = StatefulIterator(itr, start(itr))
StatefulIterator(s::StatefulIterator) = StatefulIterator(s.iter, s.state)
copy(s::StatefulIterator) = StatefulIterator(s)

eltype{T,S}(::StatefulIterator{T,S}) = eltype(T) 
eltype{T,S}(::Base.Take{StatefulIterator{T,S}}) = eltype(T)

start(s::StatefulIterator) = s
function next{T,S}(s::StatefulIterator{T,S}, ::StatefulIterator{T,S})
    (item, s.state) = next(s.iter, s.state)
    return item, s
end
done(s::StatefulIterator) = done(s.iter, s.state)
done{T,S}(s::StatefulIterator{T,S}, ::StatefulIterator{T,S}) = done(s)

read{T,S}(s::StatefulIterator{T,S}) = next(s, s)[1]
function read(s::StatefulIterator, dims::Int...)
    a = Array(eltype(s), dims)
    M = prod(dims)
    if M == 0
        return a
    end
    i = 0
    for x in s
        @inbounds a[i+=1] = x
        if i == M
            return a
        end
    end
    return a[1:i] # Should possibly be an error. Asked for too many elements.
end

function peek(s::StatefulIterator, dims::Int...)
    state = s.state
    a = read(s, dims...)
    s.state = state
    return a
end
peek(s::StatefulIterator) = next(s.iter, s.state)[1]

@generated function available{T,S}(s::StatefulIterator{T,S})
    # Returns the number of remaining bytes in stream.
    if isbits(eltype(T)) && method_exists(length, Tuple{T})
        # fast algorithm for iterators that support it
        # NB! This fails for StepRange, since state can be > length!!
        return :((length(s.iter) - s.state + 1) * sizeof(eltype(T)))
    else
        # fallback algorithm
        return quote
            state = s.state
            siz = 0
            for x in s
                siz += sizeof(x)
            end
            s.state = state
            return siz
        end
    end
end

function read{U}(s::StatefulIterator, ::Type{U})
    N = cld(sizeof(U), sizeof(eltype(s)))
    return reinterpret(U, read(s, N))[1]
end
function peek{U}(s::StatefulIterator, ::Type{U})
    N = cld(sizeof(U), sizeof(eltype(s)))
    return reinterpret(U, peek(s, N))[1]
end
function read{U}(s::StatefulIterator, ::Type{U}, dims...)
    M = prod(dims)
    N = cld(M * sizeof(U), sizeof(eltype(s)))
    return reshape(reinterpret(U, read(s, N))[1:M], dims)
end
function peek{U}(s::StatefulIterator, ::Type{U}, dims...)
    M = prod(dims)
    N = cld(M * sizeof(U), sizeof(eltype(s)))
    return reshape(reinterpret(U, peek(s, N))[1:M], dims)
end

reset!(s::StatefulIterator) = (s.state = start(s.iter))

end