Python's `asyncio` and Websockets: Building a Chat Server

Introduction
Prerequisites
Setting Up the Environment
Creating the Chat Server
Running the Chat Server
Conclusion

Introduction

In this tutorial, we will explore how to build a chat server using Python’s asyncio and websockets module. We will learn how to handle multiple connections, process incoming messages, and broadcast those messages to all connected clients. By the end of this tutorial, you will have a working chat server that can be used for real-time communication.

Prerequisites

To follow along with this tutorial, you should have a basic understanding of Python programming and web development concepts. Familiarity with asynchronous programming and the asyncio library will be helpful but not mandatory. You will also need Python 3.7 or later installed on your machine.

Setting Up the Environment

Before we begin, let’s set up our environment by creating a new directory for our project and initializing a virtual environment.

Open your terminal or command prompt.
Create a new directory for the project: mkdir chat-server
Move into the project directory: cd chat-server
Create a virtual environment: python3 -m venv venv
Activate the virtual environment:
- On macOS and Linux: source venv/bin/activate
- On Windows: venv\Scripts\activate.bat

Now that our environment is set up, let’s install the necessary dependencies.

Install the websockets library: pip install websockets

Creating the Chat Server

Initializing the Server

Let’s start by creating a new Python file called server.py in the project directory.

Create a new file called server.py.
Open server.py in your favorite text editor.

We will begin by importing the necessary libraries and creating a basic server. ```python import asyncio import websockets

# Define the server's IP address and port
HOST = "localhost"
PORT = 8000

# Define a list to store connected clients
clients = []

# Define the server's main logic
async def server(websocket, path):
    # Add the connected client to the list
    clients.append(websocket)
    print(f"New client connected: {websocket.remote_address}")

    # Wait for messages from the client
    try:
        async for message in websocket:
            print(f"Received message: {message}")
    except websockets.ConnectionClosed:
        # Remove the disconnected client from the list
        clients.remove(websocket)
        print(f"Client disconnected: {websocket.remote_address}")

# Run the server
start_server = websockets.serve(server, HOST, PORT)
asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()
``` Let's go through the code:

We import the necessary libraries: asyncio and websockets.
We define the host and port on which the server will run.
We create an empty list called clients to store connected clients.
We define the server function, which will handle each client’s connections and messages.
Inside the server function, we add the connected client to the clients list and print a message indicating a new client has connected.
We use a try...except block to wait for messages from the connected client. When a message is received, we print it to the console.
If the client’s connection is closed, we remove the client from the clients list and print a message indicating that the client has disconnected.
Finally, we run the server by creating a new websockets.serve instance and starting the event loop.

Handling Connections

Our server is now ready to accept connections and receive messages. Let’s update the code to handle incoming messages more efficiently. ```python # …

# Define the server's main logic
async def server(websocket, path):
    clients.append(websocket)
    print(f"New client connected: {websocket.remote_address}")

    try:
        async for message in websocket:
            await process_message(websocket, message)
    except websockets.ConnectionClosed:
        clients.remove(websocket)
        print(f"Client disconnected: {websocket.remote_address}")

# Define the message processing logic
async def process_message(sender, message):
    print(f"Received message from {sender.remote_address}: {message}")
    # TODO: Broadcast the message to all connected clients
    

# ...
``` We have added a new `process_message` function to handle message processing. This function will be responsible for broadcasting the received message to all connected clients. We haven't implemented the broadcasting logic yet, but we will do it shortly.

Processing Messages

Let’s update the process_message function to send the received message to all connected clients. ```python # …

# Define the message processing logic
async def process_message(sender, message):
    print(f"Received message from {sender.remote_address}: {message}")

    # Broadcast the message to all connected clients
    for client in clients:
        if client != sender:
            await client.send(f"{sender.remote_address}: {message}")

# ...
``` In the updated code, we iterate over all connected clients and check if the client is the sender of the message. If it is not the sender, we use the `client.send` method to send the message. This way, the message is broadcasted to all clients except the sender.

Sending Messages

Our server can now receive and broadcast messages. However, we haven’t written any client-side code yet. For testing purposes, let’s add a simple client that sends messages to the server.

Create a new file called client.py in the project directory.

Open client.py in your favorite text editor.

 import asyncio
 import websockets
	
 # Define the server's IP address and port
 SERVER = "ws://localhost:8000"
	
 async def main():
     async with websockets.connect(SERVER) as websocket:
         while True:
             message = input("Enter a message (or 'exit' to quit): ")
             if message == "exit":
                 break
             await websocket.send(message)
	
 asyncio.get_event_loop().run_until_complete(main())

Let’s go through the code:

We start by importing the necessary asyncio and websockets libraries.
We define the server’s address as ws://localhost:8000.
We define the main function, which is responsible for connecting to the server and sending messages.
Inside the main function, we use the websockets.connect context manager to establish a connection to the server.
We enter a loop that prompts the user to enter a message. If the message is “exit”, the loop breaks, and the program terminates. Otherwise, the message is sent to the server using websocket.send.

Disconnecting Clients

Currently, our server does not handle the disconnection event of clients. Let’s modify the code to properly remove clients when they disconnect. ```python # …

async def server(websocket, path):
    clients.append(websocket)
    print(f"New client connected: {websocket.remote_address}")

    try:
        async for message in websocket:
            await process_message(websocket, message)
    except websockets.ConnectionClosed:
        if websocket in clients:
            clients.remove(websocket)
            print(f"Client disconnected: {websocket.remote_address}")

# ...
``` We added an additional check to ensure the disconnected client is still present in the `clients` list before removing it. This prevents any potential issues if the client is already disconnected but the server still tries to remove it.

Running the Chat Server

You can now start the chat server by running the server.py file.

Open a terminal or command prompt.
Ensure you are in the chat-server project directory.
Activate the virtual environment if you haven’t done so: source venv/bin/activate (macOS/Linux) or venv\Scripts\activate.bat (Windows).
Run the server: python server.py

The server should start and display a message indicating that it is running. Keep the server running in the background.

Next, let’s start clients to test the chat functionality.

Open a new terminal or command prompt.
Ensure you are in the chat-server project directory.
Activate the virtual environment if you haven’t done so.
Run the client: python client.py

The client should connect to the server and prompt you to enter a message. Type a message and press Enter. The message should be sent to the server and broadcasted to all connected clients, including the sender.

You can run multiple instances of the client to simulate multiple users connected to the chat server.

Conclusion

In this tutorial, we’ve learned how to build a chat server using Python’s asyncio and websockets module. We covered the process of initializing the server, handling connections, processing messages, and broadcasting messages to all connected clients. We also saw how to create a simple client to send messages to the server for testing purposes.

With this knowledge, you can further enhance the chat server by implementing features such as username assignment, private messaging, or message persistence. Experiment and have fun exploring the possibilities of real-time web communication with asyncio and websockets!

Remember to stop the server by pressing Ctrl+C in the terminal running server.py.

Published: 30 July 2022