Iterator & Generator in Python

Posted on Edited on In Views: Word count in article: 285 Reading time ≈ 1 mins.
The iterator and generator are important objects in Python.

Iterator

The iterator in python is slightly different from its counterpart in C++. Once a class define __iter__(self) member function and __next__(self) member function, we can use iter() to get its iterator and next() to get the next element of this iterator:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class Number():
    def __iter__(self):
        self.n = 1
        return self

    def __next__(self):
        if self.n <= 20:
            x = self.n
            self.n += 1
            return x
        else:
            return StopIteration

n = Number()
_iter = iter(n)
print(type(_iter))
print(next(_iter))

'''
We get:
<class '__main__.Number'>
1
'''

StopIteration is the end of iterator for every iterator. It is an exception and for can catch and solve it. Therefore, we can also use for to traverse a iterator:

1
2
3
4
5
6
7
8
9
10
a = [1, 2, 3]
_iter = iter(a)
for i in _iter:
    print(i)
'''
We get:
1
2
3
'''

list, tuple, string have implemented iterators.

Generator

The generator is a special type of iterator, which makes a function iterable. Once a function uses yield to return object, this function is a generator:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
def func(n):
    print('a')
    for i in range(n):
        yield i
        print('hello')

_generator = func(10)
print(type(_generator))
print(next(_generator))
'''
We get:
<class 'generator'>
0
'''
print(next(_generator))
'''
We get:
hello
1
'''

When a function becomes a generator, it will not run immediately when called, instead, it will return a generator. When we use next() to traverse the generator, the function will run to the place where yield is, record the address and return a value. Next time, it will run starting from yield. Similarly, it will return StopIteration when finish the function. Therefore, we can also use for to traverse a generator:

1
2
for i in func(10):
    print(i + 1)