Websockets

FastAPI is a powerful, high-performance web framework that has gained popularity for its simplicity and robustness. It's built on Starlette for the web parts and Pydantic for the data parts. One of the advanced features that FastAPI provides is support for Websockets, which allows real-time communication between the client and the server.

Websockets is a communication protocol that provides full-duplex communication channels over a single TCP connection. This tutorial will guide you through setting up a WebSocket route in FastAPI.

Setting up a WebSocket Route

Before we dive into Websockets, ensure you have FastAPI installed. If not, you can install it using pip:

pip install fastapi

You also need an ASGI server, like Uvicorn or Hypercorn.

pip install uvicorn

Once you have FastAPI and an ASGI server, let's create a WebSocket route. In FastAPI, you define a WebSocket route using the websocket decorator, similar to the get or post decorators for HTTP routes.

from fastapi import FastAPI, WebSocket

app = FastAPI()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    while True:
        data = await websocket.receive_text()
        await websocket.send_text(f"Message text was: {data}")

In this example, /ws is a WebSocket route, and when a client connects to this route, it sends a message to the server. The server then sends a response back to the client with the received message.

Connecting from the Client Side

To connect to a WebSocket route from the client side, you can use the JavaScript WebSocket API:

var socket = new WebSocket("ws://localhost:8000/ws");

socket.onopen = function(e) {
  socket.send("Hello, server!");
};

socket.onmessage = function(event) {
  console.log(`[message] Data received from server: ${event.data}`);
};

socket.onclose = function(event) {
  if (event.wasClean) {
    console.log(`[close] Connection closed cleanly, code=${event.code} reason=${event.reason}`);
  } else {
    console.log('[close] Connection died');
  }
};

socket.onerror = function(error) {
  console.log(`[error] ${error.message}`);
};

Handling Disconnections

When a client disconnects, FastAPI raises a WebSocketDisconnect exception. You can catch this exception to handle disconnections gracefully:

from fastapi import FastAPI, WebSocket, WebSocketDisconnect

app = FastAPI()

class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    async def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def send_personal_message(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)
        
    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_text(message)

manager = ConnectionManager()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            await manager.send_personal_message(f"Message text was: {data}", websocket)
    except WebSocketDisconnect:
        manager.disconnect(websocket)
        await manager.broadcast(f"Client #{id(websocket)} left the chat")

In the above code, ConnectionManager handles the active WebSocket connections. When a client connects, it adds the WebSocket connection to the active_connections list. When a client disconnects, it removes the WebSocket connection from the list and sends a message to all other clients.

Websockets in FastAPI is an advanced topic that opens a whole new world of possibilities for real-time, bidirectional communication between the client and the server. It's perfect for use cases like chat applications, real-time analytics, live updates, and more.

And that's it! You now know how to use Websockets in FastAPI. Happy coding!