asyncio_state_machine.py

"""
I am trying to figure out asyncio.
I have no idea if this code is good or not, but it works.

I also have no idea how to type annotate asyncio code.
"""
from typing import List, NamedTuple
import random
import asyncio


class Node(NamedTuple):
    """
    A node has a reference to its owning state machine, as well as a list of next nodes.
    At each timestep it will choose one of the `next_nodes` at random and transition to it.
    """
    state_machine: 'StateMachine'
    next_nodes: List[int]
    name: str

    async def __call__(self, ttl: int):
        print(ttl, self.name)
        if ttl > 0:
            # Choose a random next node, and send it a smaller ttl
            next_node = self.state_machine.nodes[random.choice(self.next_nodes)]
            await next_node(ttl - 1)


class StateMachine:
    def __init__(self, transitions: List[List[int]]) -> None:
        self.nodes = [Node(self, next_nodes, f"node {i}") 
                      for i, next_nodes in enumerate(transitions)]

    async def run(self, ttl: int) -> None:  
        """
        Send the original ttl to the 0th node and hope for the best
        """          
        await self.nodes[0](ttl)


TRANSITIONS = [
    [1, 1, 2],     # 0 goes to 1 or 2
    [2],           # 1 definitely goes to 2
    [0, 1, 2, 3],  # 2 goes anywhere
    [2]            # 3 goes back to 2
]

state_machine = StateMachine(TRANSITIONS)
asyncio.run(state_machine.run(100))