Generators in Python

A generator is a function that can return several values one after another without keeping the whole set of values in memory. Let's look at it through examples:

PythonCopy

def sum(val, m):
    result = val
    while True:
        print(result)
        result += m

print(sum(0, 1))
print("Конец программы")

If you run this program, the function will loop forever, since the example provides no way to end the loop. Any code following the call to sum will never be executed.

Let's change the code:

PythonCopy

def sum(val, m):
    result = val
    while True:
        # заменили print на yield
        yield result
        result += m

print(sum(0, 1))
print("Конец программы")

# output
# >>> <generator object sum at 0x7f1032f95a50>
# >>> Конец программы

The call to sum evaluated to a <generator object>, so it did not loop endlessly.

To get concrete numbers, you need to call the generator. The returned object behaves like an iterator, and in this case it generates an infinite sequence. You can use the returned object like this:

PythonCopy

for number in sum(0, 3):
    print(number)
    if number > 11:
        break

# output
# >>> 0
# >>> 3
# >>> 6
# >>> 9
# >>> 12

Let's go back to your implementation of range: there is a class that can return the next number in a given range. In the __next__ method the program computes the current value each time and returns it to the user. Formally this can be called a generator, but Python provides a special, simplified syntax for creating one.

Let's see what the same code looks like when written using a generator.

PythonCopy

def gen_range(stop_value):
    stop_value = stop_value - 1
    current = -1
    while current < stop_value:
        current += 1
        yield current


for x in gen_range(3):
    print(x)

# output
# >>> 0
# >>> 1
# >>> 2

Generators work on the principle of remembering the function's execution context. A generator function remembers where it stopped and can resume execution after the yield keyword.

Let's look at how a generator with several yield statements behaves:

PythonCopy

def simple_generator():
    yield 1
    yield 2
    return 3

While studying iterators, we learned that you can pull values out of them manually with the next method or with a loop, and also implicitly — inside a loop. A generator offers the same capabilities: every generator is an iterator. In the previous example you used a loop; now, for variety, let's fetch the elements manually.

PythonCopy

gen = simple_generator()
print(next(gen))
print(next(gen))
print(next(gen))

As a result you will see the following:

PythonCopy

1
2
Traceback (most recent call last):
...
StopIteration: 3

So the function really does remember where it stopped after each call to next.

The question of the difference between iterators and generators in Python comes up quite often and is a relevant one. The two entities are closely related and are often confused, which leads to misunderstandings.

An iterator is a more general concept. It is an object that defines two methods: __next__ and __iter__. A generator, on the other hand, is an iterator, but not the other way around. If you look at the output of dir(<generator>), you'll find that it also defines the __next__ and __iter__ methods. It's also important to remember that you cannot get values out of a generator a second time.

Using generators

Because of their "lazy" processing, where data is handled not all at once but in parts, generators have found wide application. Let's go over the most common cases where generators can come in handy.

Infinite sequences

A generator lets you create sequences of infinite length.

PythonCopy

def func(x):
    n = x
    while True:
        n = (n + 3) * (n + 2) - 5 * n
        yield n
        if n > 100:
            break

for item in func(2):
    print(item)

# output
# >>> 10
# >>> 106

Reading large files

Working with data streams and large files is the most common use case for generators. Suppose you need to count the number of lines in a file. You can implement the task as follows:

PythonCopy

def read_file(file_name):
    with open(file_name) as f:
        return f.read().split("\n")

read_my_file = read_file("my_file.txt")

row_count = 0
for row in read_my_file:
    row_count += 1

print(f"В файле {row_count} строк.")

If the file turns out to be too large, this program will fail with an error. To avoid that, we use a generator:

PythonCopy

def read_file(file_name):
    with open(file_name) as f:
        for row in f:
            yield row