Python中实现WebSocket的示例详解_Python

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}")
    # 处理正常的消息