How can I write a higher-order generator generic over synchronousness?

user_113586 0
Viewed 156

I have a bunch of code with the same pattern:

def process(iterable):
    if inspect.isasyncgen(iterable):
        return process_async(iterable)
    else:
        return process_normal(iterable)

def process_normal(iterable):
    for i in iterable:
        do_something
        yield something

async def process_async(iterable):
    async for i in iterable:
        do_something
        yield something

This allows me to process normal and async iterables with the same code, as process() returns the same type of iterable as it gets. However, when the processing code inside the loops is non-trivial, this leads to ugly code duplication which is not always possible to extract, so I wonder, is it possible to write this more concisely, or maybe replace some of this with some stdlib things?

generator asynchronous python
2 Answers

I see red functions and I want them painted black
It’s no solution, this is just a dirty hack

Here’s a solution using a decorator. It’s something of a hack in that it takes a deep dive into how asynchronous functions and generators are actually implemented in Python; those mechanics are documented in PEP 492 and PEP 525, so they aren’t mere implementation details of CPython, but this still relies on the Python implementation never adding spurious suspension points, even if they are judged ‘harmless’, which is a little shaky assumption.

import functools, inspect

def generic_over_async(process_async):
    def process_sync(iterable, /, *args, **kwargs):
        async def iterable_async():
            for item in iterable:
                yield item
        agen = process_async(iterable_async(), *args, **kwargs)
        try:
            sent = None
            while True:
                gen = agen.asend(sent)
                try:
                    gen.send(None)
                except StopIteration as e:
                    sent = yield e.value
                else:
                    gen.throw(RuntimeError,
                        f"synchronously-called function '{process_async.__name__}' has blocked")
                finally:
                    gen.close()
        except StopAsyncIteration:
            return
        finally:
            agen.aclose()
    @functools.wraps(process_async)
    def process(iterable, /, *args, **kwargs):
        if inspect.isasyncgen(iterable):
            return process_async(iterable, *args, **kwargs)
        return process_sync(iterable, *args, **kwargs)
    return process

The above defines a decorator that takes an asynchronous generator and adds logic that:

  • when called with an asynchronous generator, it’s passed along to the decorated asynchronous generator as-is;
  • when called with a regular iterable, it’s converted to an asynchronous generator and passed to the decorated asynchronous generator. The resulting asynchronous generator is then manually driven to completion, yielding all generated values on the way.

Using this decorator successfully requires that the only thing the decorated function can await on is the iterable given and other async functions that are effectively synchronous (i.e. they never actually block); otherwise the decorator will throw a RuntimeError. Ensuring that this never happens is left as an exercise for the reader.

Test case:

import asyncio

@generic_over_async
async def process(iterable):
    async for i in iterable:
        yield i * 2

async def blow_up(iterable):
    """ Turns an iterable into an asynchronous iterable by adding dummy suspension points """
    await asyncio.sleep(0)
    for item in iterable:
        yield item
        await asyncio.sleep(0)
        
async def main():
    print(list(process(range(5))))
    print([item async for item in process(blow_up(range(5)))])

asyncio.get_event_loop().run_until_complete(main())

The above will print [0, 2, 4, 6, 8] twice.

user_38401700

You could just move processing to function and pass it as parameter in process():

def process(iterable, transform):
    if inspect.isasyncgen(iterable):
        return process_async(iterable, transform)
    else:
        return process_normal(iterable, transform)


def process_normal(iterable, transform):
    for i in iterable:
        yield transform(i)


async def process_async(iterable, transform):
    async for i in iterable:
        yield transform(i)

You could also refactor all generator into async generators and use only process_async(), but without deeper knowledge I don't know it will suit you needs.

Lastly try transforming synchronous generators into asynchronous just before processing. This will also require processing results from normal generators with asynchronous code after transformation.

async def to_async(iterable):
    for x in iterable:
        yield x


def process(iterable):
    if not inspect.isasyncgen(iterable):
        iterable = to_async(iterable)
    return process_async(iterable)


async def process_async(iterable):
    async for i in iterable:
        do_something
        yield something
user_53420200
Related