Unravelling `async for` loops

When I decided the next post in my series on Python's syntactic sugar would be on async for, I figured it would be straightforward. I have already done `for` loops, so I have something to build off of. The language reference also specifies the pseudo-code for async for, so I really didn't have to think too much about what the unravelled form should be.

But then I decided I would break with my practice of using Python 3.8 as the reference version of Python I work from and instead pull from Python 3.10 due to two additions in that version: aiter() and anext(). That decision turned out to complicate my life a bit. 😉

The bytecode

First, though, let's look at the bytecode. Our running example will be:

async for a in b:
    c

The bytecode for such a block of code is:

  2           0 LOAD_GLOBAL              0 (b)
              2 GET_AITER
        >>    4 SETUP_FINALLY           16 (to 22)
              6 GET_ANEXT
              8 LOAD_CONST               0 (None)
             10 YIELD_FROM
             12 POP_BLOCK
             14 STORE_FAST               0 (a)

  3          16 LOAD_GLOBAL              1 (c)
             18 POP_TOP
             20 JUMP_ABSOLUTE            4
        >>   22 END_ASYNC_FOR

The code for GET_AITER shows that it:

1. Calls __aiter__() on b.
2. Checks that __anext__() is defined and is awaitable on what __aiter__() returns.

The code for GET_ANEXT shows that it:

1. Checks if the iterator defines __anext__() and it is awaitable.
2. Pushes the awaitable on to the stack to be called via YIELD_FROM.

Everything else is about calling the awaitable, going through the loop, etc. It's a lot like for, but subsitute __iter__ for __aiter__ and __next__ with __anext__ which also happens to be an awaitable.

But before we talk about how to translate an async for statement to something lower level, we should talk about aiter() and anext().

The Built-ins

aiter()

If you read my for loop post and its section on iter() you will notice that when it's given an iterable it calls __iter__() on the object and checks that the iterator returned defines __next__(). I figured aiter() would be the similar and call __aiter__() and check it defined an awaitable __anext__().

Turns out I was wrong (initially). Looking at CPython 3.10.0.rc1, the implementation of aiter() calls PyObject_GetAiter(). The code itself doesn't do anything fancy:

1. Check if the iterable defines __aiter__().
2. If so, call it.
3. Check if the returned object is an async iterator via PyAiter_Check().

That's what I was expecting. But then I checked what PyAiter_Check() does, and that turned out it checked if the object defined __anext__() and __aiter__(). Huh. That's different from how iter() works.

And so I asked why this discrepancy existed. Turned out there's wasn't a key reason to it, and so for Python 3.10.0rc2, aiter() was changed to only require the returned async iterator from __aiter__() to define __anext__(), thus operating more like iter().

def aiter(iterable, /):
    """Return the async iterator for the async iterable by calling __aiter__().
    
    If the async iterator does not define an awaitable `__anext__`,
    raise `TypeError`.
    """
    iterable_type = builtins.type(iterable)
    try:
        __aiter__ = _mro_getattr(iterable_type, "__aiter__")
    except AttributeError:
        raise TypeError(f"{iterable_type.__name__!r} is not async iterable")
    else:
        iterator = __aiter__(iterable)
        iterator_type = builtins.type(iterator)
        try:
            __anext__ = _mro_getattr(iterator_type, "__anext__")
        except AttributeError:
            raise TypeError(f"{iterator_type.__name__!r} is not an async iterator")
        if not inspect.iscoroutinefunction(__anext__):
            raise TypeError(f"{iterator_type.__name__!r} is not an async iterator")
        return iterator

anext()

The implementation of anext() is straightforward as it just calls __anext__().

async def anext(iterator: AsyncIterator[Any], default: Any = _NOTHING, /) -> Any:
    """Return the next item from the async iterator by calling __anext__().

    If `default` is provided and `StopAsyncIteration` is raised, then return
    `default`.
    """
    iterator_type = builtins.type(iterator)
    try:
        __anext__ = _mro_getattr(iterator_type, "__anext__")
    except AttributeError:
        raise TypeError(f"{iterator_type.__name__!r} is not an async iterator")
    try:
        return await __anext__(iterator)
    except StopAsyncIteration:
        if default is not _NOTHING:
            return default
        else:
            raise

The semantics

Unravelling async for follows the format of for loops very closely. The only key differences are:

1. You call aiter() instead of iter()
2. You call anext() instead of next()
3. You await on anext()
4. You catch StopAsyncIteration instead of StopIteration

Thanks to having worked out how aiter() and anext() work, we can take our for loop examples and mechanically translate them to work for async for.

async for ...

_iter = aiter(b)
while True:
    try:
        a = await anext(_iter)
    except StopAsyncIteration:
        break
    else:
        c
del _iter

async for ... else ...

_iter = aiter(b)
_looping = True
while _looping:
    try:
        a = await anext(_iter)
    except StopAsyncIteration:
        _looping = False
        continue
    else:
        c
else:
    d
del _iter, _looping