Monday 12 August 2024

Python Return From Generator

In my previous post I talked about returning values (along with yielding) from generators and said that it also applied to Python. While writing the Python equivalent to the code in that post I've realised that there are more differences than I expected.

Let's start by "normal" (synchronous) iterators. The most important difference is that Python iterators yield a value in each iteration (rather than a done-value pair as JavaScript does) and throw an StopIteration exception when they reach its end (while JavaScript yields a {done: true, value: undefined}). If we return a value from the generator it's set in the value attribute of the exception. We don't have access to that exception in for-in loop, and if we do an extra next() call once outside the loop, it throws an exception again, but this time with value set to None (so this behaviour is equivalent to the JavaScript one). This means that if we want access to that value, we either use a while loop, or use a class (that as in my previous post I called AsyncGeneratorWithReturn) that wraps that loop for us. So we end up with this code (notice how in __iter2__ we leverage the yield from construct)


from typing import Any, Callable

def citiesGen():
    yield "Paris"
    yield "Porto"
    return "Europe"


print("- for of loop")
cities = citiesGen()
for city in cities:
    print(city)

try:
    print(next(cities))
except StopIteration as ex:
    print(f"return value: {ex.value}")
    
# Paris
# Porto
# None

#that's the same behaviour as in JavaScript, once it has generated the "end" (first StopIteration Exception with the value attribute set)
#the next calls raise StopIteration with value as None

print("- while loop")
#this works fine:
cities = citiesGen()
while True:
    try:
        city = next(cities)
    except StopIteration as ex:
        print(f"return value: {ex.value}")
        break


print("- GeneratorWithReturn")

class GeneratorWithReturn:
    def __init__(self, generator_fn: Callable, strategy: int):
        self.generator_fn = generator_fn
        self.result = None
        self.strategy = strategy

    # this works fine
    def __iter1__(self):
        cities = self.generator_fn()
        while True:
            try:
                yield next(cities)
            except StopIteration as ex:
                self.result = ex.value
                break

    # but it can be refactored to just one line!
    def __iter2__(self):
        self.result = yield from self.generator_fn()        


    def __iter__(self):
        return getattr(self, f"__iter{self.strategy}__")()    


for strategy in (1, 2):
    cities = GeneratorWithReturn(citiesGen, strategy)
    for city in cities:
        print(city)
    print(f"return value: {cities.result}")


I said in my previous post that this idea of yielding and finally returning a value seemed mainly interesting to me for async iteration. We know that same as JavaScript Python features async generators (and async for loops). And the interesting thing is that they come with a limitation, we can not return values from them (we get a #SyntaxError: 'return' with value in async generator) and we can not use yield from (the equivalent to JavaScript's yield *). However, this is not a big deal, we can easily workaround the limitation by yielding an object of a particular class (that I've called GeneratorReturn) that wraps the return value, and wrapping the async for in an AsyncGeneratorWithReturn class that takes care of this return value (rather than taking care of a StopAsyncIteration exception).



import asyncio

class GeneratorReturn:
    def __init__(self, value):
        self.value = value

class AsyncGeneratorWithReturn:
    def __init__(self, agen_fn):
        self.agen_fn = agen_fn
        self.result = None

    async def __aiter__(self):
        agen = self.agen_fn()
        while True:
            if isinstance(res := await anext(agen), GeneratorReturn):
                self.result = res.value
                return
            else:
                yield res


async def designVacations():
    destinations = []
    await asyncio.sleep(1)
    destinations.append("Paris")
    yield "destination found: Paris"
    await asyncio.sleep(1)
    destinations.append("Porto")
    yield "destination found: Porto"
    	#return (" -> ").join(destinations)
    	#SyntaxError: 'return' with value in async generator
    yield GeneratorReturn((" -> ").join(destinations))


async def async_main():
    print("- AsyncGeneratorWithReturn");
    vacationsProcess = AsyncGeneratorWithReturn(designVacations)
    async for city in vacationsProcess:
        print(city)
    print(f"return value (travelPlan): ${vacationsProcess.result}")

    # destination found: Paris
    # destination found: Porto
    # return value (travelPlan): $Paris -> Porto

if __name__ == "__main__":
    asyncio.run(async_main())


It's interesting to notice that while Javascript for await can be used both with asynchronous and synchonous iterators, Python async for only works with asynchronous iterators. In Javascript the iteration method of asynchronous and synchronous iterators is called next, while in Python we have 2 different methods: anext() and next(). All this is related to the fact that while in JavaScript we can await for a Promise and for a normal value, in Python we can only await for an awaitable (coroutine, Future or Task).

No comments:

Post a Comment