Lua-style coroutines

lua_coroutines.py

"""
Lua-style coroutines.

Examples taken from https://www.lua.org/pil/9.html.
"""

import types


class Coroutine:
    SUSPENDED = "suspended"
    RUNNING = "running"
    DEAD = "dead"

    @classmethod
    def coroutine(cls, func):
        func = types.coroutine(func)

        @types.coroutine
        def _coroutine():
            args = yield
            return (yield from func(*args))

        def _start():
            co = _coroutine()
            co.send(None)
            return cls(co)

        return _start

    def __init__(self, coroutine):
        self.coroutine = coroutine
        self.status = self.SUSPENDED

    def resume(self, *args):
        self.status = self.RUNNING
        try:
            value = self.coroutine.send(args)
        except StopIteration as e:
            self.status = self.DEAD
            return e.value

        self.status = self.SUSPENDED
        return value

    def __repr__(self):
        return f"<Coroutine status={self.status!r}>"


if __name__ == '__main__':
    @Coroutine.coroutine
    def coroutine(a, b, c):
        print("co", a, b, c)
        yield 1
        yield 2
        return "Done"

    co = coroutine()
    print(co.resume(1, 2, 3), co)
    print(co.resume(4), co)
    print(co.resume(5), co)

    # A call to resume returns, after the true that signals no errors, any
    # arguments passed to the corresponding yield:
    @Coroutine.coroutine
    def coroutine(a, b):
        yield a + b, a - b

    co = coroutine()
    print(co.resume(20, 10))

    # Symmetrically, yield returns any extra arguments passed to the
    # corresponding resume:
    @Coroutine.coroutine
    def coroutine():
        print("co", (yield))

    co = coroutine()
    co.resume()
    co.resume(4, 5)

    # Finally, when a coroutine ends, any values returned by its main function
    # go to the corresponding resume:
    @Coroutine.coroutine
    def coroutine():
        yield from []  # Hack to force this to be a generator
        return 6, 7

    co = coroutine()
    print(co.resume())

    # Producer/Consumer example
    def receive(prod: Coroutine):
        return prod.resume()

    @types.coroutine
    def send(x):
        yield x

    @Coroutine.coroutine
    async def producer():
        while True:
            x = input("? ")
            await send(x)

    def filter(prod):
        @Coroutine.coroutine
        async def _filter():
            line = 1
            while True:
                x = receive(prod)
                await send("%5d %s" % (line, x))
                line = line + 1

        return _filter()

    def consumer(prod):
        while True:
            x = receive(prod)
            print(x)

    # Simple pipeline
    consumer(filter(producer()))