Python 的 WebSocket 实现详解

💖 欢迎来到我的博客！ 非常高兴能在这里与您相遇。在这里，您不仅能获得有趣的技术分享，还能感受到轻松愉快的氛围。无论您是编程新手，还是资深开发者，都能在这里找到属于您的知识宝藏，学习和成长。

🔍 博客内容包括：

• Java核心技术与微服务：涵盖Java基础、JVM、并发编程、Redis、Kafka、Spring等，帮助您全面掌握企业级开发技术。
• 大数据技术：涵盖Hadoop（HDFS）、Hive、Spark、Flink、Kafka、Redis、ECharts、Zookeeper等相关技术。
• 开发工具：分享常用开发工具（IDEA、Git、Mac、Alfred、Typora等）的使用技巧，提升开发效率。
• 数据库与优化：总结MySQL及其他常用数据库技术，解决实际工作中的数据库问题。
• Python与大数据：专注于Python编程语言的深度学习，数据分析工具（如Pandas、NumPy）和大数据处理技术，帮助您掌握数据分析、数据挖掘、机器学习等技术。
• 数据结构与算法：总结数据结构与算法的核心知识，提升编程思维，帮助您应对大厂面试挑战。

🌟 我的目标：持续学习与总结，分享技术心得与解决方案，和您一起探索技术的无限可能！在这里，我希望能与您共同进步，互相激励，成为更好的自己。

📣 欢迎订阅本专栏，与我一起在这个知识的海洋中不断学习、分享和成长！💻🚀

---

📍版权声明：本博客所有内容均为原创，遵循CC 4.0 BY-SA协议，转载请注明出处。

目录

[1. WebSocket 基础概念](#1. WebSocket 基础概念)

[2. Python WebSocket 库](#2. Python WebSocket 库)

[3. 安装 websockets 库](#3. 安装 websockets 库)

[4. WebSocket 服务器端实现](#4. WebSocket 服务器端实现)

[4.1. 创建 WebSocket 服务器](#4.1. 创建 WebSocket 服务器)

[4.2. 代码说明](#4.2. 代码说明)

[5. WebSocket 客户端实现](#5. WebSocket 客户端实现)

[5.1. 创建 WebSocket 客户端](#5.1. 创建 WebSocket 客户端)

[5.2. 代码说明](#5.2. 代码说明)

[6. 异常处理](#6. 异常处理)

[6.1. 客户端异常处理](#6.1. 客户端异常处理)

[6.2. 服务器端异常处理](#6.2. 服务器端异常处理)

[7. WebSocket 高级使用](#7. WebSocket 高级使用)

[7.1. 广播消息](#7.1. 广播消息)

[7.2. 认证和授权](#7.2. 认证和授权)

[8. 总结](#8. 总结)

---

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

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

1. WebSocket 基础概念

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

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

WebSocket 握手过程：

1. 客户端向服务器发起一个 HTTP 请求，并要求升级到 WebSocket 协议。
2. 服务器响应客户端的请求，确认升级协议。
3. 客户端和服务器建立 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 服务器端实现

4.1. 创建 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())

4.2. 代码说明

• handle_connection 是处理每个客户端连接的异步函数。每当有客户端连接时，服务器将接收到的消息打印出来，并将消息返回给客户端。
• websockets.serve() 用于启动 WebSocket 服务器，监听指定的地址和端口（此处为 localhost 和 8765）。
• await server.wait_closed() 保证服务器持续运行，直到手动关闭。

5. WebSocket 客户端实现

5.1. 创建 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())

5.2. 代码说明

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

8. 总结

本文详细介绍了如何在 Python 中使用 websockets 库实现 WebSocket 服务端和客户端，涉及到的主要概念包括：

• WebSocket 协议基础及其优势；
• 如何实现 WebSocket 服务器和客户端；
• 如何处理异常、进行消息广播和认证等高级应用。