WebSocket Server and Client in Python

WebSocket is a protocol providing full-duplex communication channels over a single TCP connection. It’s often used for real-time applications such as chat applications, live notifications, etc.

Installing the websockets library

First, ensure you have the websockets library installed. You can install it using pip:

pip install websockets

WebSocket Server

Let’s start by creating a WebSocket server.

import asyncio
import websockets

connected_clients = set()

async def handle_client(websocket, path):
    # Add the new client connection to the set
    connected_clients.add(websocket)
    try:
        async for message in websocket:
            print(f"Received message: {message}")
            # Broadcast the message to all connected clients
            for client in connected_clients:
                if client != websocket:
                    await client.send(message)
    except websockets.exceptions.ConnectionClosed as e:
        print(f"Client disconnected: {e}")
    finally:
        # Remove the client from the set when it disconnects
        connected_clients.remove(websocket)

start_server = websockets.serve(handle_client, 'localhost', 8765)

asyncio.get_event_loop().run_until_complete(start_server)
print("WebSocket server started on ws://localhost:8765")
asyncio.get_event_loop().run_forever()

Explanation

Imports:

import asyncio
import websockets

• asyncio: Python library for writing concurrent code using the async/await syntax.
• websockets: A library for building WebSocket servers and clients in Python.

Connected Clients Set:

connected_clients = set()

• connected_clients: A set to keep track of all connected clients.

Handle Client Function:

async def handle_client(websocket, path):
    connected_clients.add(websocket)
    try:
        async for message in websocket:
            print(f"Received message: {message}")
            for client in connected_clients:
                if client != websocket:
                    await client.send(message)
    except websockets.exceptions.ConnectionClosed as e:
        print(f"Client disconnected: {e}")
    finally:
        connected_clients.remove(websocket)

• handle_client: An asynchronous function to handle communication with a client.
• async for message in websocket: Listens for messages from the client.
• Broadcasting: Sends the received message to all other connected clients.

Start Server:

start_server = websockets.serve(handle_client, 'localhost', 8765)
asyncio.get_event_loop().run_until_complete(start_server)
print("WebSocket server started on ws://localhost:8765")
asyncio.get_event_loop().run_forever()

• websockets.serve: Creates the WebSocket server.
• run_until_complete: Runs the server until it completes (which it never does in this case).
• run_forever: Keeps the server running indefinitely.

WebSocket Client

Now, let’s create a WebSocket client to connect to this server.

import asyncio
import websockets

async def listen():
    uri = "ws://localhost:8765"
    async with websockets.connect(uri) as websocket:
        async for message in websocket:
            print(f"Message from server: {message}")

async def send():
    uri = "ws://localhost:8765"
    async with websockets.connect(uri) as websocket:
        while True:
            message = input("Enter message to send: ")
            await websocket.send(message)

async def main():
    await asyncio.gather(listen(), send())

asyncio.get_event_loop().run_until_complete(main())

Explanation

Imports:

import asyncio
import websockets

• asyncio: Python library for writing concurrent code using the async/await syntax.
• websockets: A library for building WebSocket servers and clients in Python.

Listening for Messages:

async def listen():
    uri = "ws://localhost:8765"
    async with websockets.connect(uri) as websocket:
        async for message in websocket:
            print(f"Message from server: {message}")

• listen: An asynchronous function to listen for messages from the server.
• async for message in websocket: Continuously listens for messages from the server.

Sending Messages:

async def send():
    uri = "ws://localhost:8765"
    async with websockets.connect(uri) as websocket:
        while True:
            message = input("Enter message to send: ")
            await websocket.send(message)

• send: An asynchronous function to send messages to the server.
• await websocket.send(message): Sends the input message to the server.

Main Function:

async def main():
    await asyncio.gather(listen(), send())

• main: Gathers both the listen and send functions to run them concurrently.

Running the Client:

asyncio.get_event_loop().run_until_complete(main())

• Runs the main function until it completes (which it never does in this case).

Running the Server and Client

Start the Server:

python websocket_server.py

Run the Client:

python websocket_client.py