make data struct iterable in python

Question

Hey everyone,

So I've done some looking online and throught the documentation, but I'm having troubling finding out how to do this. I am working on creating an adventure game. I have a level class (which contains a bunch of rooms) and a Room class. I would like to be able to do something like the following.

l = Level()
for room in l:
    #following code here

This seems pretty simple to me, however I can't find the correct procedures for implementing a custom iterator. This may be an easy answer, but I'm having trouble finding it online. Any help is greatly appreciated!

Answer 1

Use __iter__ with an iterator-generator. E.g.

def __iter__(self):
  for r in rooms:
    yield r

An iterator-generator is basically a psuedo-method used to implement an iterator. Note that there is no requirement that the generator use a for loop. It can use any combination of constructs (if, for, while, etc.) as needed. Basically, you just have to remember that the caller will get elements in the order you "call" yield, and the iteration will end when the method does.

See this section of the Python tutorial.

Answer 2

If you define a member method called __iter__(self), Python will know how to iterate through it.

In this method, I suggest you use yield to return your data (this is called a generator). You could also return a list or tuple, but this is more efficient memory-wise. Here is an example:

class Test(object):
    def __iter__(self):
        for x in range(10): 
            yield x

l = Test()
for room in l:
    print room

Answer 3

As an alternative to writing a generator, the __iter__ method just needs to return an iterator - if your Level object has an internal datastructure that holds the rooms then you could return it's iterator directly:

class Level(object):
    def __init__(self):
        self.rooms = []

    def __iter__(self):
        return iter(self.rooms)

If the rooms container is a dictionary e.g. mapping room names to Room objects then you can get a iterator to the Room objects with the dict.itervalues method:

class Level(object):
    def __init__(self):
        self.rooms = {}

    def __iter__(self):
        return self.rooms.itervalues()

Answer 4

Why not use generator.

>>> import timeit
>>> class Test:
    def __iter__(self):
        for i in range(10):
            yield i

>>> t1 = lambda: [k for k in Test()]
>>> timeit.timeit(t1)
3.529460948082189


>>> def test2():
    for i in range(10):
        yield i

>>> t2 = lambda: [k for k in test2()]
>>> timeit.timeit(t2)
3.171831107365392