How to make a custom object iterable?

Posted on

Question :

How to make a custom object iterable?

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 #1:

__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.
Answered By: mgilson

Answer #2:

simply implementing __iter__ should be enough.

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

    def __iter__(self):
        #return iter(self.__stations[1:]) #uncomment this if you wanted to skip the first element.
        return iter(self.__stations)


a = direction(1)
a._direction__stations= range(5)

b = direction(1)
b._direction__stations = range(10)

import itertools
print list(itertools.chain.from_iterable([a,b]))
print list(itertools.chain.from_iterable([range(5),range(10)]))

output:

[0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
[0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

See here for why it’s _direction__stations

Any identifier of the form __spam (at least two leading underscores,
at most one trailing underscore) is textually replaced with
classname_spam, where classname is the current class name with leading underscore(s) stripped.

Answered By: M4rtini

Answer #3:

You can subclass list as well:

class Direction(list):
    def __init__(self, seq=[], id_=None):
        list.__init__(self,seq)
        self.id = id_ if id_ else id(self)

    def __iter__(self):
        it=list.__iter__(self) 
        next(it)                       # skip the first...
        return it  

d=Direction(range(10))
print(d)       # all the data, no iteration
# [0, 1, 2, 3, 4]

print (', '.join(str(e) for e in d))     # 'for e in d' is an iterator
# 1, 2, 3, 4

ie, skips the first.

Works for nested lists as well:

>>> d1=Direction([range(5), range(10,15), range(20,25)])
>>> d1
[range(0, 5), range(10, 15), range(20, 25)]
print(list(itertools.chain.from_iterable(d1)))
[10, 11, 12, 13, 14, 20, 21, 22, 23, 24]          
Answered By: dawg

Leave a Reply

Your email address will not be published. Required fields are marked *