Denver Dang
Denver Dang

Reputation: 2615

Restart a for loop after break

I have an iterator that consists of several lists of the same size. For my purpose I need to know the length of at least one of these lists. But as it is with iterators they can't be accessed the same way as ordinary arrays. So my idea was to get this length by saying:

for i in iter:
    list_len = len(i)
    break

And this works, however, when using this list later on, and wanting to loop over it again it skips the first iteration, and basically continues from the next iteration from the previous loop (the one above).

Is there some way to fix this ? Or, what is the pythonic way of doing it ? I was thinking/reading about doing it like:

from itertools import tee

iter_tmp, iter = tee(iter)
for i in iter_tmp:
    list_len = len(i)
    break

And yeah, that works too, since I can now use the original iter for later use, but it just hurt my eyes that I have to make a loop, import itertools and such just to get the length of a list in an iterator. But maybe that is just the way to go about it ?

UPDATE

Just trying to further explain what I'm doing.

As such iterations is not a list or an array, but in my case, if I were to loop through my iterator I would get something like (in the case of my iterator having four "lists" in it):

>>> for i in iter_list:
        print(i)
[1, 2, 5, 3]
[3, 2, 5, 8]
[6, 8, 3, 7]
[1, 4, 6, 1]

Now, all "lists" in the iterator has the same length, but since the lists themselves are calculated through many steps, I really don't know the length in any way before it enters the iterator. If I don't use an iterator I run out of memory - so it is a pro/con solution. But yeah, it is the length of just one of the lists I need as a constant I can use throughout the rest of my code.

Upvotes: 1

Views: 940

Answers (3)

Moberg
Moberg

Reputation: 5503

That is how iterators work. But you have a few options apart from tee.

You can extract the first element and reuse it when iterating the second time:

first_elem = next(my_iter)
list_len = len(first_elem)

for l in itertools.chain([first_elem], my_iter):
    pass

Or if you are going to iterate over the iterator more times, you could perhaps listify it (if it's feasible to fit in memory).

my_list = list(my_iter)
first_len = len(my_list[0])

for l in my_list:
    pass

And certainly not the least, as Palivek said, keep/get the information about the length of the lists (from) somewhere else.

Upvotes: 2

code_onkel
code_onkel

Reputation: 2947

You can do a little trick and wrap the original iterator in a generator. This way, you can obtain the first element and "re-yield" it with the generator without consuming the entire iterator. The head() function below returns the first element and a generator that iterates over the original sequence.

def head(seq):
    seq_iter = iter(seq)
    first = next(seq_iter)
    def gen():
        yield first
        yield from seq_iter
    return first, gen()

seq = range(100, 300, 50)
first, seq2 = head(seq)

print('first item: {}'.format(first))

for item in seq2:
    print(item)

Output:

first item: 100
100
100
150
200
250

This is conceptually equivalent to Moberg's answer, but uses a generator to "re-assemble" the original sequence instead of itertools.chain().

Upvotes: 0

bipll
bipll

Reputation: 11940

In general iterators are not re-iteratable so you'll probably need to store something additional anyway.

class peek_iterator(object):
    def __init__(self, source):
        self._source = iter(source)
        self._first = None
        self._sent = False

    def __iter__(self):
        return self

    def next(self):
        if self._first is None:
            self._first = self._source.next()
        if self._sent:
            return self._source.next()
        self._sent = True
        return self._first

    def get_isotropic(self, getter):
        if self._first is None:
            self._first = self._source.next()
        return getter(self._first)

lists = [[1, 2, 3], [4, 5, 6]]

i = peek_iterator(lists)

print i.get_isotropic(len) # 3

for j in i: print j        # [1, 2, 3]; [4, 5, 6]

Upvotes: 0

Related Questions