chat.py

import asyncio
from colorama import Fore, Back, Style
from random import choice

def random_color():
    return choice([Fore.BLACK, Fore.RED, Fore.GREEN, Fore.YELLOW, Fore.BLUE, Fore.MAGENTA, Fore.CYAN, Fore.WHITE])


class Pool:

    def __init__(self):
        self.connections = set()

    def broadcast(self, writer, msg):
        for user in self.connections:
            # We don't want the writer to see the message twice
            if user != writer:
                user.write(f"{msg}\n".encode())

    def broadcast_new_message(self, writer, msg):
        self.broadcast(writer, f"[{writer.color + writer.nick + Style.RESET_ALL}]: {msg}")

    def add_connection(self, writer):
        self.connections.add(writer)

    def remove_connection(self, writer):
        self.connections.remove(writer)


pool = Pool()


async def handle_connection(reader, writer):
    writer.write("Welcome, please choose a nickname: ".encode())
    nick = await reader.readuntil("\n".encode())
    writer.nick = nick.decode().strip()
    writer.color = random_color()

    # Now we add the user to the pool
    pool.add_connection(writer)
    pool.broadcast(writer, f"{nick.decode().strip()} just joined")

    while True:
        try:
            # Read the next message from this connected user (writer)
            data = await reader.readuntil(b"\n")
        except (asyncio.exceptions.IncompleteReadError, ConnectionError):
            break

        msg = data.decode().strip()
        pool.broadcast_new_message(writer, msg)

        # Wait for write to finish
        await writer.drain()

    # Cleaning up the connection...
    writer.close()
    await writer.wait_closed()
    pool.remove_connection(writer)

async def main():
    server = await asyncio.start_server(handle_connection, "0.0.0.0", 8888)

    async with server:
        await server.serve_forever()


asyncio.run(main())