第31天：WebSocket实时通信

# day31_websocket_server.py
# 功能：WebSocket服务器实现
# 需要安装：pip install websockets

import asyncio
import websockets
import json
from datetime import datetime

# 存储所有连接的客户端
connected_clients = set()


async def handle_client(websocket, path):
    """
    处理客户端连接

    参数：
        websocket: WebSocket连接对象
        path: 连接路径
    """
    # 添加到连接集合
    connected_clients.add(websocket)
    client_id = id(websocket)

    print(f"[{datetime.now()}] 客户端连接: {client_id}")
    print(f"当前在线: {len(connected_clients)} 人")

    # 发送欢迎消息
    welcome_msg = {
        "type": "system",
        "message": f"欢迎加入聊天室！当前在线 {len(connected_clients)} 人",
        "timestamp": datetime.now().isoformat()
    }
    await websocket.send(json.dumps(welcome_msg, ensure_ascii=False))

    # 广播用户加入消息
    join_msg = {
        "type": "user_join",
        "message": f"用户 {client_id} 加入了聊天室",
        "online_count": len(connected_clients),
        "timestamp": datetime.now().isoformat()
    }
    await broadcast(json.dumps(join_msg, ensure_ascii=False), exclude=websocket)

    try:
        # 持续接收消息
        async for message in websocket:
            print(f"[{datetime.now()}] 收到消息: {message}")

            try:
                # 解析JSON消息
                data = json.loads(message)

                # 处理不同类型的消息
                if data.get('type') == 'ping':
                    # 心跳响应
                    pong_msg = {
                        "type": "pong",
                        "timestamp": datetime.now().isoformat()
                    }
                    await websocket.send(json.dumps(pong_msg))

                elif data.get('type') == 'chat':
                    # 聊天消息，广播给所有人
                    chat_msg = {
                        "type": "chat",
                        "from": client_id,
                        "message": data.get('message', ''),
                        "timestamp": datetime.now().isoformat()
                    }
                    await broadcast(json.dumps(chat_msg, ensure_ascii=False))

            except json.JSONDecodeError:
                # 非JSON消息，直接广播
                text_msg = {
                    "type": "chat",
                    "from": client_id,
                    "message": message,
                    "timestamp": datetime.now().isoformat()
                }
                await broadcast(json.dumps(text_msg, ensure_ascii=False))

    except websockets.exceptions.ConnectionClosed:
        print(f"[{datetime.now()}] 客户端 {client_id} 连接关闭")

    finally:
        # 移除连接
        connected_clients.remove(websocket)
        print(f"当前在线: {len(connected_clients)} 人")

        # 广播用户离开消息
        leave_msg = {
            "type": "user_leave",
            "message": f"用户 {client_id} 离开了聊天室",
            "online_count": len(connected_clients),
            "timestamp": datetime.now().isoformat()
        }
        await broadcast(json.dumps(leave_msg, ensure_ascii=False))


async def broadcast(message, exclude=None):
    """
    广播消息给所有客户端

    参数：
        message: 要广播的消息
        exclude: 排除的WebSocket连接（可选）
    """
    if connected_clients:
        # 过滤掉exclude的连接
        clients = connected_clients - {exclude} if exclude else connected_clients

        # 并发发送给所有客户端
        await asyncio.gather(
            *[client.send(message) for client in clients],
            return_exceptions=True
        )


async def main():
    """启动WebSocket服务器"""
    print("""
    ╔══════════════════════════════════════════════════════════╗
    ║       WebSocket 服务器                                    ║
    ║       WebSocket Server                                   ║
    ╚══════════════════════════════════════════════════════════╝

    服务器地址: ws://localhost:8765
    等待客户端连接...
    """)

    # 启动服务器
    server = await websockets.serve(
        handle_client,
        "localhost",
        8765
    )

    print("WebSocket服务器已启动")
    print("按 Ctrl+C 停止服务器\n")

    # 保持服务器运行
    await asyncio.Future()  # 永久运行


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n服务器已停止")

# day31_websocket_client.py
# 功能：WebSocket客户端

import asyncio
import websockets
import json
from datetime import datetime


async def client():
    """WebSocket客户端"""
    uri = "ws://localhost:8765"

    print("""
    ╔══════════════════════════════════════════════════════════╗
    ║       WebSocket 客户端                                    ║
    ║       WebSocket Client                                   ║
    ╚══════════════════════════════════════════════════════════╝
    """)

    async with websockets.connect(uri) as websocket:
        print(f"已连接到服务器: {uri}\n")

        # 启动接收消息的任务
        receive_task = asyncio.create_task(receive_messages(websocket))

        # 启动心跳任务
        heartbeat_task = asyncio.create_task(send_heartbeat(websocket))

        # 发送消息
        print("输入消息并按Enter发送（输入'quit'退出）：\n")
        try:
            while True:
                message = await asyncio.get_event_loop().run_in_executor(
                    None, input, "你: "
                )

                if message.lower() == 'quit':
                    break

                # 发送聊天消息
                chat_msg = {
                    "type": "chat",
                    "message": message,
                    "timestamp": datetime.now().isoformat()
                }
                await websocket.send(json.dumps(chat_msg, ensure_ascii=False))

        except KeyboardInterrupt:
            pass
        finally:
            receive_task.cancel()
            heartbeat_task.cancel()
            print("\n已断开连接")


async def receive_messages(websocket):
    """接收服务器消息"""
    try:
        async for message in websocket:
            data = json.loads(message)

            if data['type'] == 'system':
                print(f"\n[系统] {data['message']}")
            elif data['type'] == 'user_join':
                print(f"\n[通知] {data['message']} (在线: {data['online_count']})")
            elif data['type'] == 'user_leave':
                print(f"\n[通知] {data['message']} (在线: {data['online_count']})")
            elif data['type'] == 'chat':
                print(f"\n[{data['from']}]: {data['message']}")
            elif data['type'] == 'pong':
                # 心跳响应，不显示
                pass

            print("你: ", end='', flush=True)

    except asyncio.CancelledError:
        pass


async def send_heartbeat(websocket):
    """发送心跳"""
    try:
        while True:
            await asyncio.sleep(30)  # 每30秒
            ping_msg = {
                "type": "ping",
                "timestamp": datetime.now().isoformat()
            }
            await websocket.send(json.dumps(ping_msg))
    except asyncio.CancelledError:
        pass


if __name__ == "__main__":
    asyncio.run(client())