pydantic、websockets和asyncio三者同时使用

pydantic

pydantic 是一个Python库，用于数据解析和验证。它允许你定义数据的结构，并在数据被赋值给变量时自动进行类型检查和数据验证。这对于构建RESTful API、消息队列消费者等场景非常有用，因为它可以确保数据的准确性和一致性。

websockets

websockets 是一个用于Python的WebSocket客户端和服务器库。WebSocket是一种在单个长时间连接上提供全双工通信渠道的协议。这使得客户端和服务器可以实时地交换数据，非常适合需要实时更新（如聊天应用、实时数据可视化等）的应用。

asyncio

asyncio是Python标准库的一部分，它提供了实现协程、事件循环、任务、异步I/O等功能的工具，使得Python开发者可以更方便地编写异步代码，提高程序的并发性能和响应性

当pydantic、websockets和asyncio三者同时使用时，你可以构建一个基于WebSocket的实时API，其中pydantic用于数据验证，websockets处理WebSocket通信，而asyncio则用于异步处理。以下是一个简化的例子，展示了这三者如何协同工作。

场景

假设我们正在构建一个实时聊天应用。客户端通过WebSocket连接到服务器，并发送包含用户名和消息的数据。服务器验证这些数据，然后将消息广播给所有连接的客户端。

代码示例

1. 定义消息模型（使用Pydantic）

from pydantic import BaseModel, ValidationError  
from datetime import datetime  
  
class ChatMessage(BaseModel):  
    username: str  
    message: str  
    timestamp: datetime = datetime.now()  # 自动设置时间戳  
  
    class Config:  
        # 允许通过字段名赋值（可选）  
        allow_population_by_field_name = True

2. WebSocket服务器（使用Websockets和Asyncio）

import asyncio  
import websockets  
from pydantic import ValidationError  
from chat_message import ChatMessage  
  
# 存储所有连接的WebSocket  
connected = set()  
  
async def chat_handler(websocket, path):  
    connected.add(websocket)  
    try:  
        async for message in websocket:  
            try:  
                # 尝试将接收到的消息解析为ChatMessage对象  
                data = ChatMessage.parse_raw(message)  
                print(f"Received from {data.username}: {data.message}")  
  
                # 将消息广播给所有连接的客户端（除了发送者）  
                for ws in connected:  
                    if ws != websocket:  
                        await ws.send(f"{data.username}: {data.message}")  
  
            except ValidationError as e:  
                print(f"Invalid message: {e}")  
  
    finally:  
        connected.remove(websocket)  
  
async def main():  
    async with websockets.serve(chat_handler, "localhost", 6789):  
        await asyncio.Future()  # 运行直到被取消  
  
# 运行服务器  
asyncio.run(main())

注意：上面的代码示例中，ChatMessage.parse_raw(message) 并不是 pydantic 的标准方法。实际上，你应该先将JSON字符串解码为Python字典，然后再使用 ChatMessage.parse_obj(data) 来验证和解析数据。这里为了简化，我们假设message已经是正确的JSON格式，并且直接传递给parse_raw（这实际上需要你自定义或调整代码）。

• 在生产环境中，你需要处理更多的错误情况，比如网络中断、无效的WebSocket连接等。
• 广播消息时，你可能需要更复杂的逻辑来确保消息只发送给应该接收它们的客户端（例如，基于聊天室的分组）。
• 使用asyncio.Future()来保持服务器运行是一个简单的技巧，但在生产环境中，你可能需要一个更优雅的关闭机制。
• 考虑到性能和可扩展性，你可能需要使用更高级的WebSocket库或框架，如FastAPI与Uvicorn（支持ASGI）结合websockets，或者使用Django Channels等。

3. 客户端

虽然这个示例主要关注后端，但客户端通常会是一个Web应用或移动应用，它使用WebSocket API连接到服务器，并发送和接收消息。客户端代码将依赖于你选择的平台和技术栈。

这里提供一个使用Python编写的WebSocket客户端的示例。这个客户端将连接到我们之前定义的WebSocket服务器，并发送聊天消息。我们将使用websockets库来创建客户端连接。

首先，确保你已经安装了websockets库。如果没有，你可以通过pip安装它：

接下来是客户端的Python代码示例：

import asyncio  
import json  
import websockets  
  
async def chat_client(uri):  
    async with websockets.connect(uri) as websocket:  
        while True:  
            # 等待用户输入用户名和消息  
            username = input("Enter your username: ")  
            message = input("Enter your message: ")  
  
            # 将消息和用户名编码为JSON  
            data = json.dumps({"username": username, "message": message})  
  
            # 发送消息到服务器  
            await websocket.send(data)  
  
            # 等待并打印服务器的响应  
            response = await websocket.recv()  
            print(f"Received: {response}")  
  
            # 为了演示，我们只发送一条消息然后退出循环（在真实应用中，你可能会放在一个不同的循环中）  
            break  # 移除这行以持续发送消息  
  
# WebSocket服务器的URI  
uri = "ws://localhost:6789"  
  
# 运行客户端  
asyncio.run(chat_client(uri))

但是，请注意，上面的代码有几个问题：

1. 它会在发送一条消息后退出循环（由于break语句）。在真实的聊天应用中，你可能希望将输入放在一个单独的循环中，以便用户可以连续发送多条消息。
2. 它直接将用户名和消息编码为JSON字符串，但服务器端的ChatMessage.parse_raw并不存在。实际上，你应该在服务器端使用ChatMessage.parse_obj(json.loads(message))来解析JSON。

下面是一个改进后的客户端示例，它允许用户连续发送消息，并且假设服务器端已经正确设置来解析JSON消息：

import asyncio  
import json  
import websockets  
  
async def chat_client(uri):  
    async with websockets.connect(uri) as websocket:  
        while True:  
            # 等待用户输入用户名和消息  
            username = input("Enter your username (or type 'exit' to quit): ")  
            if username.lower() == 'exit':  
                break  
  
            message = input("Enter your message: ")  
  
            # 将消息和用户名编码为JSON  
            data = json.dumps({"username": username, "message": message})  
  
            # 发送消息到服务器  
            await websocket.send(data)  
  
            # 等待并打印服务器的响应（如果有的话）  
            try:  
                response = await websocket.recv_str()  # 使用recv_str()如果服务器发送的是纯文本  
                print(f"Received: {response}")  
            except websockets.ConnectionClosed:  
                print("The server closed the connection.")  
                break  
  
# WebSocket服务器的URI  
uri = "ws://localhost:6789"  
  
# 运行客户端  
asyncio.run(chat_client(uri))

请注意，在await websocket.recv()上使用了recv_str()，这是因为websockets库默认返回bytes对象，但在这个例子中，我们假设服务器发送的是UTF-8编码的字符串。如果你的服务器发送的是JSON字节串，你可能需要先使用json.loads()来解码它。然而，在这个例子中，我们假设服务器已经将JSON解码为字符串并发送了回来。

另外，请确保你的服务器端代码能够处理多个并发连接，并且正确地将消息广播给所有连接的客户端（除了发送者）。在上面的服务器端示例中，我已经添加了将新连接的WebSocket添加到connected集合中的逻辑，但你需要确保在发送消息时不会尝试向已经关闭的WebSocket发送数据。