asyncio echo server

asyncio_echo.py

import random

import asyncio
import colored


host = "127.0.0.1"
port = 8080


async def _echo_handler(reader, writer):
    while True:
        msg = await reader.read(1024)

        if not msg:
            writer.close()
            await writer.wait_closed()
            return

        writer.write(msg)
        await writer.drain()


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


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


async def run_client(_id):
    reader, writer = await asyncio.open_connection(host, port)

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

    msg = _mk_msg(_id)
    writer.write(msg.encode())

    data = await reader.read(1024)
    print(msg)
    print(_render_response(data))

    writer.close()
    await writer.wait_closed()


async def _main():
    server = await asyncio.start_server(_echo_handler, host, port)

    async with server:
        clients = tuple(run_client(i) for i in range(5))

        await asyncio.gather(*clients)

        server.close()


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