Compiling Generators · python/mypy Wiki · GitHub
Skip to content

Compiling Generators

Jump to bottom
Jukka Lehtosalo edited this page Oct 25, 2022 · 3 revisions

Generator functions must maintain state between invocations. A generator function is compiled into two classes. One is an environment class that maintains the local variables of the function, and the label/location where execution will continue (__mypyc_next_label__ attribute). The second class implements the actual generator. The environment class is similar to what mypyc uses to support nested functions.

For example, consider this simple generator:

def f(x: int) -> Iterator[int]:
    print('a')
    yield x
    print('b')

This is compiled to something roughly like this (in heavily simplified pseudo-python):

def f(x: int) -> Iterator[int]:
    _env = f_env()
    _env.x = x
    _env.__mypyc_next_label__ = 0
    _gen = f_gen()
    _gen.__mypyc_env__ = _env
    return _gen

class f_env:
    x: int
    __mypyc_next_label__: int

class f_gen:
    __mypyc_env__: f_env

    def __next__(self) -> int:
        _env = self.__mypyc_env__
        _label = _env.__mypyc_next_label__
        if _label == 0:
            print('a')
            _ret = _env.x
            _env.__mypyc_next_label__ = 1
            return _ret
        elif _label == 1:
            print('b')
            _env.__mypyc_next_label__ = -1
            raise StopIteration()
        else:
            raise StopIteration()  # Already completed

The actual generated code is more complicated because of error handling, and close/throw/send methods.

Clone this wiki locally