Diego Aguayo
Diego Aguayo

Reputation: 33

Iterate twice through an object

I'm trying to create an iterable object, and when I do 1 loop it is okay, but when doing multiple loops, it doesn't work. Here is my simplified code:

class test():

    def __init__(self):
        self.n = 0

    def __iter__(self):
        return self

    def __next__(self):
        if self.n < len(self)-1:
            self.n += 1
            return self.n
        else:
            raise StopIteration

    def __len__(self):
        return 5

#this is an example iteration
test = test()
for i in test:
    for j in test:
        print(i,j)

#it prints is 
1 2
1 3
1 4

#What i expect is
1 1 
1 2
1 3
1 4
2 1
2 2
2 3
...
4 3
4 4

How can I make this object (in this case test) to iterate twice and get all the combinations of number i and j in the example loop?

Upvotes: 2

Views: 596

Answers (2)

chepner
chepner

Reputation: 530920

You want an instance of test to be iterable, but not its own iterator. What's the difference?

An iterable is something that, upon request, can supply an iterator. Lists are iterable, because iter([1,2,3]) returns a new listiterator object (not the list itself). To make test iterable, you just need to supply an __iter__ method (more on how to define it in a bit).

An iterator is something that, upon request, can produce a new element. It does this by calling its __next__ method. An iterator can be thought of as two pieces of information: a sequence of items to produce, and a cursor indicating how far along that sequence it currently is. When it reaches the end of its sequence, it raises a StopIteration exception to indicate that the iteration is at an end. To make an instance an iterator, you supply a __next__ method in its class. An iterator should also have a __iter__ method that just returns itself.

So how do you make test iterable without being an iterator? By having its __iter__ method return a new iterator each time it is called, and getting rid of its __next__ method. The simplest way to do that is to make __iter__ a generator function. Define your class something like:

class Test():

    def __init__(self):
        self._size = 5

    def __iter__(self):
        n = 0
        while n < self._size:
            yield n
            n += 1

    def __len__(self):
        return self._size

Now when you write

test = Test()
for i in test:    # implicit call to iter(test)
    for j in test:  # implicit call to iter(test)
        print(i, j)

i and j both draw values from separate iterators over the same iterable. Each call to test.__iter__ returns a different generator object that keeps track of its own n.

Upvotes: 3

Sunny Patel
Sunny Patel

Reputation: 8078

Take a look at itertools.product.

You should be able to accomplish what you're looking for:

from itertools import product
...
test = test()
for i, j in product(test, repeat=2):
    print(i,j)

I love this library!

Upvotes: 0

Related Questions