How to make a custom object iterable?

Matthieu Riegler picture Matthieu Riegler · Feb 9, 2014 · Viewed 25k times · Source

I have a list of custom-class objects (sample is below).

Using: list(itertools.chain.from_iterable(myBigList)) I wanted to "merge" all of the stations sublists into one big list. So I thought I need to make my custom class an iterable.

Here is a sample of my custom class.

class direction(object) :
    def __init__(self, id) :
        self.id = id              
        self.__stations = list()

    def __iter__(self):
        self.__i = 0                #  iterable current item 
        return iter(self.__stations)

    def __next__(self):
        if self.__i<len(self.__stations)-1:
            self.__i += 1         
            return self.__stations[self.__i]
        else:
            raise StopIteration

I implemented __iter__ and __next__ but it doesn't seems to work. They're not even called.

Any idea what could I've done wrong?

Note: Using Python 3.3

Answer

mgilson picture mgilson · Feb 9, 2014

__iter__ is what gets called when you try to iterate over a class instance:

>>> class Foo(object):
...     def __iter__(self):
...         return (x for x in range(4))
... 
>>> list(Foo())
[0, 1, 2, 3]

__next__ is what gets called on the object which is returned from __iter__ (on python2.x, it's next, not __next__ -- I generally alias them both so that the code will work with either...):

class Bar(object):
   def __init__(self):
       self.idx = 0
       self.data = range(4)
   def __iter__(self):
       return self
   def __next__(self):
       self.idx += 1
       try:
           return self.data[self.idx-1]
       except IndexError:
           self.idx = 0
           raise StopIteration  # Done iterating.
   next = __next__  # python2.x compatibility.

In the comments, it was asked how you would construct and object that could be iterated multiple times. In this case, I'd recommend taking the same approach that Python takes and split the iterator from the data container:

class BarIterator(object):
   def __init__(self, data_sequence):
       self.idx = 0
       self.data = data_sequence
   def __iter__(self):
       return self
   def __next__(self):
       self.idx += 1
       try:
           return self.data[self.idx-1]
       except IndexError:
           self.idx = 0
           raise StopIteration  # Done iterating.


class Bar(object):
    def __init__(self, data_sequence):
        self.data_sequence = data_sequence
    def __iter__(self):
        return BarIterator(self.data_sequence)