How to build a basic iterator?

Viewed 430362

How would one create an iterative function (or iterator object) in python?

10 Answers

Iterator objects in python conform to the iterator protocol, which basically means they provide two methods: __iter__() and __next__().

  • The __iter__ returns the iterator object and is implicitly called at the start of loops.

  • The __next__() method returns the next value and is implicitly called at each loop increment. This method raises a StopIteration exception when there are no more value to return, which is implicitly captured by looping constructs to stop iterating.

Here's a simple example of a counter:

class Counter:
    def __init__(self, low, high):
        self.current = low - 1
        self.high = high

    def __iter__(self):
        return self

    def __next__(self): # Python 2: def next(self)
        self.current += 1
        if self.current < self.high:
            return self.current
        raise StopIteration


for c in Counter(3, 9):
    print(c)

This will print:

3
4
5
6
7
8

This is easier to write using a generator, as covered in a previous answer:

def counter(low, high):
    current = low
    while current < high:
        yield current
        current += 1

for c in counter(3, 9):
    print(c)

The printed output will be the same. Under the hood, the generator object supports the iterator protocol and does something roughly similar to the class Counter.

David Mertz's article, Iterators and Simple Generators, is a pretty good introduction.

First of all the itertools module is incredibly useful for all sorts of cases in which an iterator would be useful, but here is all you need to create an iterator in python:

yield

Isn't that cool? Yield can be used to replace a normal return in a function. It returns the object just the same, but instead of destroying state and exiting, it saves state for when you want to execute the next iteration. Here is an example of it in action pulled directly from the itertools function list:

def count(n=0):
    while True:
        yield n
        n += 1

As stated in the functions description (it's the count() function from the itertools module...) , it produces an iterator that returns consecutive integers starting with n.

Generator expressions are a whole other can of worms (awesome worms!). They may be used in place of a List Comprehension to save memory (list comprehensions create a list in memory that is destroyed after use if not assigned to a variable, but generator expressions can create a Generator Object... which is a fancy way of saying Iterator). Here is an example of a generator expression definition:

gen = (n for n in xrange(0,11))

This is very similar to our iterator definition above except the full range is predetermined to be between 0 and 10.

I just found xrange() (suprised I hadn't seen it before...) and added it to the above example. xrange() is an iterable version of range() which has the advantage of not prebuilding the list. It would be very useful if you had a giant corpus of data to iterate over and only had so much memory to do it in.

If you looking for something short and simple, maybe it will be enough for you:

class A(object):
    def __init__(self, l):
        self.data = l

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

example of usage:

In [3]: a = A([2,3,4])

In [4]: [i for i in a]
Out[4]: [2, 3, 4]

All answers on this page are really great for a complex object. But for those containing builtin iterable types as attributes, like str, list, set or dict, or any implementation of collections.Iterable, you can omit certain things in your class.

class Test(object):
    def __init__(self, string):
        self.string = string

    def __iter__(self):
        # since your string is already iterable
        return (ch for ch in self.string)
        # or simply
        return self.string.__iter__()
        # also
        return iter(self.string)

It can be used like:

for x in Test("abcde"):
    print(x)

# prints
# a
# b
# c
# d
# e

Include the following code in your class code.

 def __iter__(self):
        for x in self.iterable:
            yield x

Make sure that you replace self.iterablewith the iterable which you iterate through.

Here's an example code

class someClass:
    def __init__(self,list):
        self.list = list
    def __iter__(self):
        for x in self.list:
            yield x


var = someClass([1,2,3,4,5])
for num in var: 
    print(num) 

Output

1
2
3
4
5

Note: Since strings are also iterable, they can also be used as an argument for the class

foo = someClass("Python")
for x in foo:
    print(x)

Output

P
y
t
h
o
n
class uc_iter():
    def __init__(self):
        self.value = 0
    def __iter__(self):
        return self
    def __next__(self):
        next_value = self.value
        self.value += 2
        return next_value

Improving previous answer, one of the advantage of using class is that you can add __call__ to return self.value or even next_value.

class uc_iter():
    def __init__(self):
        self.value = 0
    def __iter__(self):
        return self
    def __next__(self):
        next_value = self.value
        self.value += 2
        return next_value
    def __call__(self):
        next_value = self.value
        self.value += 2
        return next_value
c = uc_iter()
print([c() for _ in range(10)])
print([next(c) for _ in range(5)])
# [0, 2, 4, 6, 8, 10, 12, 14, 16, 18]
# [20, 22, 24, 26, 28]

Other example of a class based on Python Random that can be both called and iterated could be seen on my implementation here

Related