python websockets，如何设置连接超时

import asyncio
import websockets

async def connect_with_timeout(uri, timeout):
    try:
        # 创建一个协程，用于建立WebSocket连接
        connect_coroutine = websockets.connect(uri)
        
        # 使用asyncio.wait_for函数设置连接超时
        connection = await asyncio.wait_for(connect_coroutine, timeout)
        
        # 连接成功，进行后续操作
        await do_something_with_connection(connection)
        
    except asyncio.TimeoutError:
        # 连接超时，进行异常处理
        print("连接超时")

async def do_something_with_connection(connection):
    # 在连接上进行后续操作，例如发送和接收数据
    await connection.send("Hello, WebSocket!")
    response = await connection.recv()
    print(response)

# 设置WebSocket服务器的URI和连接超时时间
uri = "ws://example.com/websocket"
timeout = 5  # 5秒

# 运行连接函数
asyncio.run(connect_with_timeout(uri, timeout))