Saturday, February 14, 2009

Custom Python iterators.

I've been experimenting with writing custom Python iterators. In Python, any sequence type can be iterated over. For example, when you say "for i in list_obj", list_obj actually returns an iterator. For this common case, the developer need not be concerned with the iterator. They know that it will always behave correctly in a for loop.

If you have a class that needs to be iterated over, you can define a custom iterator for your class as demonstrated below.
#Example; Custom iterators.

class MyIterator:
def __init__(self, obj):
self.obj=obj
self.cnt=0

def __iter__(self):
return self

def next(self):
try:
result=self.obj.get(self.cnt)
self.cnt+=1
return result
except IndexError:
raise StopIteration

class MyClass:
def __init__(self):
self.data=["Item 1", "Item 2", "Item 3"]

def __iter__(self):
return MyIterator(self)

def get(self, index):
print "Getting item..."
return self.data[index]

if __name__=="__main__":
for item in MyClass():
print item
Here, we have a custom iterator defined, MyIterator. The constructor initializes the object it will iterate over as well as the counter. The __iter__() method is required to return the current MyIterator instance so the iterator can also be iterated over. Sort of like meta-iteration.

The important method is next() which returns the next item in the iteration. In this case, every time next() is invoked, we attempt to return the value of get() and increment the counter. Once there are no more items, we raise a StopIteration exception, which will gracefully exit the iteration.

Finally, you'll notice that MyClass defines a __iter__() method. This is what gives instances of MyClass that ability to be iterated over. The method simply returns a MyIterator instance.