Python中实现WebSocket的案例代码-QQ沐编程

WebSocket 是一种计算机通信协议，它为客户端和服务器之间的双向通信提供了一个全双工的通道。通过 WebSocket，客户端和服务器可以在一个长期的连接中保持通信，而无需每次都建立新的连接。这个特性在实时聊天、在线游戏、股票行情等应用中非常有用。

本篇文章将详解 Python 中如何实现 WebSocket，内容包括 WebSocket 的基础概念、Python 中的 WebSocket 库、简单的客户端和服务器端代码示例、以及一些高级使用场景。

1. WebSocket 基础概念

WebSocket 协议最初是由 IETF（Internet Engineering Task Force）发布的，目的是提供一个更高效的通信协议，可以通过单个长连接在客户端和服务器之间进行双向通信。WebSocket 的通信通过在 HTTP 协议上建立一个持久连接（即 WebSocket 握手），并且支持全双工通信，即数据可以同时从客户端发送到服务器，也可以从服务器发送到客户端。

WebSocket 与传统的 HTTP 协议有很大的不同。HTTP 是一个请求-响应式的协议，客户端每次需要获取数据时都需要向服务器发送请求，而 WebSocket 则是通过一个初始握手请求将通信通道建立起来，之后双方可以随时发送数据。

WebSocket 握手过程：

客户端向服务器发起一个 HTTP 请求，并要求升级到 WebSocket 协议。
服务器响应客户端的请求，确认升级协议。
客户端和服务器建立 WebSocket 连接，可以开始全双工通信。

2. Python WebSocket 库

在 Python 中，我们通常使用以下几种库来实现 WebSocket：

websockets：一个轻量级的 Python WebSocket 实现，支持异步编程。
socket.io：一个更为高级的库，支持事件驱动和自动重连等特性，常用于实时应用。
asyncio：虽然 asyncio 本身并不是一个 WebSocket 库，但它为 WebSocket 的异步编程提供了基础支持。

在本篇文章中，我们将主要介绍 websockets 库，并使用它来实现 WebSocket 客户端和服务器。

3. 安装 websockets 库

在使用 websockets 库之前，我们需要先安装它。可以通过 pip 来安装：

pip install websockets

4. WebSocket 服务器端实现

创建 WebSocket 服务器

WebSocket 服务器通常运行在一个特定的端口，并等待客户端的连接请求。我们可以通过 websockets 库来实现一个简单的 WebSocket 服务器。

import asyncio
import websockets
 
async def handle_connection(websocket, path):
    # 接收到客户端的消息
    print(f"Connected to {websocket.remote_address}")
    try:
        async for message in websocket:
            print(f"Received message: {message}")
            # 发送消息回客户端
            await websocket.send(f"Server received: {message}")
    except Exception as e:
        print(f"Error: {e}")
    finally:
        print(f"Disconnected from {websocket.remote_address}")
 
async def main():
    # 启动 WebSocket 服务器，监听 8765 端口
    server = await websockets.serve(handle_connection, "localhost", 8765)
    print("WebSocket server started on ws://localhost:8765")
    await server.wait_closed()
 
# 启动事件循环
asyncio.run(main())

代码说明

handle_connection 是处理每个客户端连接的异步函数。每当有客户端连接时，服务器将接收到的消息打印出来，并将消息返回给客户端。

websockets.serve() 用于启动 WebSocket 服务器，监听指定的地址和端口（此处为 localhost 和 8765）。

await server.wait_closed() 保证服务器持续运行，直到手动关闭。

5. WebSocket 客户端实现

创建 WebSocket 客户端

WebSocket 客户端用于与服务器建立连接，并发送和接收消息。我们使用 websockets.connect 来连接到服务器。

import asyncio
import websockets
 
async def send_message():
    uri = "ws://localhost:8765"
    async with websockets.connect(uri) as websocket:
        # 向服务器发送消息
        await websocket.send("Hello, Server!")
        print("Message sent to server.")
        
        # 接收服务器返回的消息
        response = await websocket.recv()
        print(f"Received from server: {response}")
 
# 启动事件循环
asyncio.run(send_message())

代码说明

websockets.connect(uri) 用于与 WebSocket 服务器建立连接。uri 是服务器的地址，此处为 ws://localhost:8765。

使用 await websocket.send() 向服务器发送消息。

使用 await websocket.recv() 接收来自服务器的消息。

6. 异常处理

在实际应用中，WebSocket 连接可能会因为网络中断、服务器关闭等原因而失败。我们可以在客户端和服务器端都加入适当的异常处理。

6.1 客户端异常处理

async def send_message():
    uri = "ws://localhost:8765"
    try:
        async with websockets.connect(uri) as websocket:
            await websocket.send("Hello, Server!")
            print("Message sent to server.")
            response = await websocket.recv()
            print(f"Received from server: {response}")
    except websockets.exceptions.ConnectionClosed as e:
        print(f"Connection closed: {e}")
    except Exception as e:
        print(f"Error: {e}")
 
asyncio.run(send_message())

6.2 服务器端异常处理

async def handle_connection(websocket, path):
    print(f"Connected to {websocket.remote_address}")
    try:
        async for message in websocket:
            print(f"Received message: {message}")
            await websocket.send(f"Server received: {message}")
    except websockets.exceptions.ConnectionClosed as e:
        print(f"Connection closed: {e}")
    except Exception as e:
        print(f"Error: {e}")
    finally:
        print(f"Disconnected from {websocket.remote_address}")

7. WebSocket 高级使用

7.1. 广播消息

如果你想让服务器向所有已连接的客户端广播消息，可以维护一个客户端列表并发送消息。

clients = set()
 
async def handle_connection(websocket, path):
    clients.add(websocket)
    try:
        async for message in websocket:
            print(f"Received message: {message}")
            # 广播消息给所有客户端
            for client in clients:
                if client != websocket:
                    await client.send(f"Broadcast message: {message}")
    except Exception as e:
        print(f"Error: {e}")
    finally:
        clients.remove(websocket)
        print(f"Disconnected from {websocket.remote_address}")

7.2. 认证和授权

你可以使用 WebSocket 协议的初始握手阶段进行认证和授权，比如使用 HTTP header 传递认证信息。

async def handle_connection(websocket, path):
    headers = websocket.request_headers
    token = headers.get('Authorization')
    if token != "expected_token":
        await websocket.send("Unauthorized")
        await websocket.close()
        return
    print(f"Client authorized with token: {token}")
    # 处理正常的消息