Python coroutine based state machine

I needed a simple state machine in a project I'm working on. Doing some research I came across coroutines. This post won't even try to explain what coroutines are. Instead, it just shows a super simple example of a finite state machine using coroutine.



This example is based on this and this. If you want to learn more about python coroutine, here is a great post.

def coroutine(func):
    """
    coroutine decorator, to automatically call next() upon creation
    """
    def start(*args,**kwargs):
        cr = func(*args,**kwargs)
        cr.next()
        return cr
    return start


@coroutine
def StateMachine():
    """
    The state machine
    """
    state = 'A'
    while True:
        print(state)
        input = (yield)
        if state == 'A':
            if input == 'a':
                state = 'B'
            elif input == 'b':
                state = 'C'
            else:
                break
        elif state == 'B':
            if input == 'a':
                state = 'C'
            elif input == 'b':
                state = 'A'
            else:
                break
        elif state == 'C':
            if input == 'a':
                state = 'A'
            elif input == 'b':
                state = 'B'
            else:
                break
        else:
            break

if __name__ == '__main__':
    sm = StateMachine()
    sm.send('a')
    sm.send('a')
    sm.send('b')
    sm.send('a')
    sm.send('b')
    sm.send('b')
    sm.send('b')
    sm.send('a')
    # sm.send('oops')