How to implement the Python's `itertools.tee` function

Internals & Design
1

Python’s doc has given a roughly equivalent to its itertools.tee function:

def tee(iterable, n=2):
    it = iter(iterable)
    deques = [collections.deque() for i in range(n)]
    def gen(mydeque):
        while True:
            if not mydeque:             # when the local deque is empty
                try:
                    newval = next(it)   # fetch a new value and
                except StopIteration:
                    return
                for d in deques:        # load it to all the deques
                    d.append(newval)
            yield mydeque.popleft()
    return tuple(gen(d) for d in deques)

However, I do not know how to implement the yield keyword. Is there an equivalent way of doing that?

2

The most direct translation would be to use Tasks.

3

There’s an old pre-1.0 blog post on making iterators using tasks, but I think the “new” (it’s now been a while) iterators interface might make this easier/more idiomatic.

4

Julia’s yield/yieldto is closer to Python’s await, not Python’s yield (which is more “light weight” as it does not need a task scheduler). For non-IO and non-parallel iterators, I don’t think it’s a good idea to use tasks/channels in Julia.

To implement tee in Julia, you need to do something like

function tee(iter, n=2)
    deques = [[] for _ in 1:n]
    return [Tee(iter, mydeque, deques) for mydeque in deques]
end

struct Tee{T}
    iter::T
    mydeque::Vector
    deques::Vector{Vector}
end

Base.iterate(iter::Tee, state = nothing) =
    ...

See https://docs.julialang.org/en/latest/manual/interfaces/#man-interface-iteration-1 for how to define Base.iterate.

Another approach is to use @yield from ResumableFunctions.jl.

(Advanced topic: Above approach uses Vector{Any} as mydeque. Using boxed value like this not very optimal and it’s better to use mutate-or-widen approach. But I don’t think it’s possible to do this for deques.)

3 Likes