复制代码
#!/usr/bin/env python
import asyncio
import json
import time
from websockets.asyncio.server import serve
async def stream_data(websocket):
"""流式发送数据的处理函数"""
# 接收客户端的初始消息
try:
message = await websocket.recv()
print(f"<<< 收到客户端消息: {message}")
# 解析客户端请求
try:
request = json.loads(message)
stream_type = request.get('type', 'text')
count = request.get('count', 10)
except json.JSONDecodeError:
stream_type = 'text'
count = 10
print(f"开始流式发送 {count} 条 {stream_type} 数据...")
# 根据不同类型流式发送数据
if stream_type == 'text':
# 文本流式传输
for i in range(count):
response = {
'type': 'text',
'data': f'这是第 {i+1} 条文本消息',
'timestamp': time.time(),
'progress': f'{i+1}/{count}'
}
await websocket.send(json.dumps(response, ensure_ascii=False))
print(f">>> 发送: {response['data']}")
await asyncio.sleep(1) # 模拟延迟
elif stream_type == 'number':
# 数字流式传输
for i in range(count):
response = {
'type': 'number',
'data': i + 1,
'square': (i + 1) ** 2,
'timestamp': time.time(),
'progress': f'{i+1}/{count}'
}
await websocket.send(json.dumps(response))
print(f">>> 发送数字: {response['data']}")
await asyncio.sleep(0.5)
elif stream_type == 'json':
# JSON数据流式传输
sample_data = [
{'id': 1, 'name': '张三', 'age': 25, 'city': '北京'},
{'id': 2, 'name': '李四', 'age': 30, 'city': '上海'},
{'id': 3, 'name': '王五', 'age': 28, 'city': '广州'},
{'id': 4, 'name': '赵六', 'age': 35, 'city': '深圳'},
{'id': 5, 'name': '钱七', 'age': 22, 'city': '杭州'}
]
for i, item in enumerate(sample_data[:count]):
response = {
'type': 'json',
'data': item,
'index': i,
'total': len(sample_data[:count]),
'progress': f'{i+1}/{len(sample_data[:count])}'
}
await websocket.send(json.dumps(response, ensure_ascii=False))
print(f">>> 发送JSON数据: {item}")
await asyncio.sleep(1.5)
# 发送结束信号
end_message = {
'type': 'end',
'message': '数据传输完成',
'total_sent': count,
'timestamp': time.time()
}
await websocket.send(json.dumps(end_message, ensure_ascii=False))
print(">>> 数据传输完成")
except Exception as e:
print(f"处理客户端连接时出错: {e}")
error_response = {
'type': 'error',
'message': str(e),
'timestamp': time.time()
}
await websocket.send(json.dumps(error_response, ensure_ascii=False))
async def main():
print("WebSocket流式服务器启动中...")
print("支持的流式类型: text(文本), number(数字), json(JSON数据)")
print("请使用客户端连接 localhost:8765")
async with serve(stream_data, "localhost", 8765) as server:
await server.serve_forever()
if __name__ == "__main__":
try:
asyncio.run(main())
except KeyboardInterrupt:
print("\n服务器已停止")
复制代码
#!/usr/bin/env python
import asyncio
import json
import websockets
# 全局变量存储WebSocket连接
_global_websocket = None
_global_uri = "ws://localhost:8765"
async def get_websocket_connection():
"""获取全局WebSocket连接，只连接一次"""
global _global_websocket
if _global_websocket is None:
print("正在建立WebSocket连接...")
_global_websocket = await websockets.connect(_global_uri)
print(f"已连接到服务器: {_global_uri}")
return _global_websocket
async def close_websocket_connection():
"""关闭全局WebSocket连接"""
global _global_websocket
if _global_websocket is not None:
await _global_websocket.close()
_global_websocket = None
print("WebSocket连接已关闭")
async def reset_websocket_connection():
"""重置全局WebSocket连接"""
global _global_websocket
_global_websocket = None
async def test_stream_client(stream_type="text", count=5):
"""测试流式数据接收的客户端"""
try:
# 获取全局连接（只连接一次）
websocket = await get_websocket_connection()
# 发送请求配置
request = {
"type": stream_type,
"count": count
}
await websocket.send(json.dumps(request))
print(f"已发送请求: {request}")
# 接收流式数据
received_count = 0
while True:
try:
message = await websocket.recv()
data = json.loads(message)
if data['type'] == 'end':
print(f"\n✅ 传输完成! 总共接收 {data.get('total_sent', 0)} 条数据")
break
elif data['type'] == 'error':
print(f"\n❌ 错误: {data['message']}")
break
else:
received_count += 1
print(f"\n📥 第 {received_count} 条数据:")
print(f" 类型: {data['type']}")
print(f" 内容: {data['data']}")
print(f" 进度: {data.get('progress', 'N/A')}")
if 'timestamp' in data:
print(f" 时间戳: {data['timestamp']}")
except websockets.exceptions.ConnectionClosed:
print("连接已关闭")
# 重置全局连接
await reset_websocket_connection()
break
except Exception as e:
print(f"连接错误: {e}")
# 出错时重置连接
await reset_websocket_connection()
async def interactive_client():
"""交互式客户端"""
print("=== WebSocket流式数据客户端 ===")
print("支持的类型: text(文本), number(数字), json(JSON数据)")
print("注意: WebSocket连接只会建立一次")
try:
while True:
stream_type = "text"
if stream_type == 'quit':
break
count_input = input("请输入要接收的数据条数 (默认5): ").strip()
count = int(count_input) if count_input.isdigit() else 5
print(f"\n开始接收 {count} 条 {stream_type} 类型的数据...")
await test_stream_client(stream_type, count)
print("继续执行")
except KeyboardInterrupt:
print("\n客户端已退出")
except Exception as e:
print(f"发生错误: {e}")
finally:
# 程序结束时关闭连接
await close_websocket_connection()
if __name__ == "__main__":
# 可以选择运行交互式客户端或直接测试
# choice = input("选择模式 - 1: 交互式, 2: 直接测试text类型 (默认1): ").strip()
#
# if choice == "2":
# # 直接测试
# asyncio.run(test_stream_client("text", 3))
# else:
# # 交互式模式
# asyncio.run(interactive_client())
asyncio.run(interactive_client())