Skip to Content
Course content

281: WebSockets in Python with websockets

Click on the "Edit" button in the top corner of the screen to edit your slide content.

I've spent a lot of time building APIs, and for the longest time, I tried to force everything into a Request-Response pattern. But recently, I wanted to build a dashboard that shows my server's CPU and memory usage in real-time. If I used a standard REST API, I'd have to poll the server every second. It's wasteful, it creates a lot of overhead, and it feels... wrong. That's where WebSockets come in. They keep a persistent connection open, allowing the server to push data to the client whenever it wants.

Why my first attempt failed

I started by installing the websockets library and tried to write a simple script to broadcast system stats. My first instinct was to treat it like a regular Python function. I wrote a loop to send data, but as soon as I ran it, I got a TypeError telling me that the function was a coroutine and hadn't been awaited. I forgot that websockets is built entirely on asyncio.

import asyncio
import websockets
import psutil # I'm using psutil for the system stats

async def send_stats(websocket):
    # I thought I could just loop here
    while True:
        cpu = psutil.cpu_percent()
        await websocket.send(f"CPU Load: {cpu}%")
        await asyncio.sleep(1)

async def main():
    async with websockets.serve(send_stats, "localhost", 8765):
        await asyncio.Future() # This keeps the server running forever

asyncio.run(main())

This worked, but only for one person. The moment I opened a second browser tab to check the stats, the first tab stopped receiving updates or the server behaved unpredictably. I realized my send_stats function was handling the connection, but the way I structured the main loop was a bit too simplistic for a production-ready feel. More importantly, I noticed that if I closed the browser tab, my terminal started screaming with ConnectionClosedError exceptions.

Cleaning up the crash

It's a classic mistake: assuming the client will always be there. In a WebSocket world, clients vanish constantly. They refresh the page, they lose Wi-Fi, or they just close the tab. If you don't handle that, your server loop will crash the moment a user leaves.

I wrapped the send loop in a try...except block. I also decided to change the data format. Sending raw strings is fine for a quick test, but if I want to send both CPU and RAM, JSON is the way to go. It makes the client-side parsing much cleaner.

import asyncio
import websockets
import psutil
import json

async def monitor_client(websocket):
    try:
        while True:
            # Gathering multiple stats into a dictionary
            stats = {
                "cpu": psutil.cpu_percent(),
                "ram": psutil.virtual_memory().percent
            }
            await websocket.send(json.dumps(stats))
            await asyncio.sleep(1)
    except websockets.exceptions.ConnectionClosedOK:
        print("Client disconnected gracefully")
    except websockets.exceptions.ConnectionClosedError:
        print("Client disconnected abruptly")

async def main():
    # 'serve' creates a server that calls 'monitor_client' for every new connection
    async with websockets.serve(monitor_client, "localhost", 8765):
        print("Server started on ws://localhost:8765")
        await asyncio.Future() 

asyncio.run(main())

Testing the handshake

Now, you can't test this with a browser's address bar because that sends an HTTP request, not a WebSocket upgrade request. I usually use a browser's DevTools console for a quick test. If you open a browser tab, hit F12, and go to the Console, you can run this:

const socket = new WebSocket('ws://localhost:8765');
socket.onmessage = function(event) {
    console.log('Data from server:', JSON.parse(event.data));
};

I watched the console and saw the CPU and RAM percentages ticking every second. The "magic" here is the handshake. The client asks the server, "Hey, can we upgrade this connection to a WebSocket?" The server says "Sure," and from that point on, the TCP connection stays open. No more headers being sent back and forth every second, just raw data frames flowing through the pipe.

Considering the "broadcast" problem

One thing I noticed as I played with this: my current setup is 1-to-1. Each client gets their own loop. If I wanted to send a global alert to 100 connected users simultaneously, I wouldn't want 100 different loops calculating the same CPU percentage. I'd need a way to track all active connections in a set() and iterate through them.

That's a bit more advanced, but it's the logical next step. You'd create a global CONNECTED_CLIENTS = set(), add the websocket object to it when the function starts, and remove it in the finally block. Then, you can have one separate task that broadcasts to everyone at once.




📋 Practical Task

Build a Real-Time System Log Streamer

Your task is to create a WebSocket server that streams "log" messages to any connected client. Instead of system stats, this server should simulate a live log file.

  • Create a server using websockets.serve on port 8766.
  • Inside the handler, create a loop that sends a randomly generated log message (e.g., "INFO: User logged in", "ERROR: Database timeout", "WARN: High memory usage") every 2 seconds.
  • Ensure the messages are sent as JSON objects containing both a timestamp (use datetime.now().isoformat()) and the message.
  • Implement error handling so that the server doesn't crash when a client disconnects.
  • Print a message to the server console whenever a new client connects and whenever one disconnects.
Rating
0 0

There are no comments for now.

to be the first to leave a comment.