Redis list iterator

list_iter.py

def lrange_iter(r, key, N_max=None, step=10**5):
    if N_max is None:
        starts = itertools.count(step=step)
    else:
        starts = range(0, N_max, step)

    for start in starts:
        end = start + step - 1
        if N_max is not None:
            end = min(end, N_max - 1)
        vals = r.lrange(key, start, end)
        if vals:
            yield from vals
        else:
            break


def destructive_lrange_iter(r, key, N_max=None, step=10**5):
    def ends():
        if N_max is None:
            yield from itertools.repeat(step)
        else:
            for start in range(0, N_max, step):
                yield min(step, N_max - start)

    def multi_lpop(key, N):
        vals = r.lrange(key, 0, N-1)
        if vals:
            r.ltrim(key, len(vals), -1)
        return vals

    for N in ends():
        vals = multi_lpop(key, N)
        if vals:
            yield from vals
        else:
            break