trio echo

trio_echo.py

import colored
import random
import trio

host = "127.0.0.1"
port = 8080


def _mk_msg(_id):
    return f"query from {_id}"


async def client_sender(client_stream, _id):
    data = _mk_msg(_id)
    print(data)
    await client_stream.send_all(data.encode())


def _render_response(msg, color="green"):
    if isinstance(msg, bytes):
        msg = msg.decode()
    return f"> {colored.fg(color)}{msg}{colored.attr('reset')}"


async def client_receiver(client_stream):
    data = await client_stream.receive_some(1024)
    if not data:
        return
    print(_render_response(data))


async def client(_id):
    client_stream = await trio.open_tcp_stream(host, port)

    await trio.sleep(random.randint(0, 5))

    async with client_stream:
        async with trio.open_nursery() as n:
            n.start_soon(client_sender, client_stream, _id)
            n.start_soon(client_receiver, client_stream)


async def _echo_handler(server_stream):
    while True:
        data = await server_stream.receive_some(1024)
        if not data:
            return
        await server_stream.send_all(data)


async def _main():
    with trio.fail_after(2) as timeout_scope:
        async with trio.open_nursery() as outer:
            await outer.start(trio.serve_tcp, _echo_handler, port)

            async with trio.open_nursery() as n:
                for _id in range(5):
                    n.start_soon(client, _id)

            outer.cancel_scope.cancel()
    print("try syncing from another peer!")


if __name__ == "__main__":
    trio.run(_main)