Compiling Generators

Generator functions must maintain state between invocations. A generator function is compiled into two classes. One is an environment class that maintains local variables of the function. The second one is the actual generator. The environment class is similar to 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')
            raise StopIteration()

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