119 lines
3.2 KiB
Python
119 lines
3.2 KiB
Python
import asyncio
|
|
import logging
|
|
from threading import Thread
|
|
|
|
from flandre.nodes.Node import Node
|
|
from flandre.utils.Msg import ImageArgMsg
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
import uvicorn
|
|
|
|
from typing import List
|
|
|
|
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, APIRouter
|
|
|
|
app = FastAPI()
|
|
|
|
|
|
# class ConnectionManager:
|
|
# def __init__(self):
|
|
# # 存放激活的ws连接对象
|
|
# self.active_connections: List[WebSocket] = []
|
|
# self.router = APIRouter()
|
|
# self.router.add_api_route("/hello", self.hello, methods=["GET"])
|
|
#
|
|
#
|
|
# def hello(self):
|
|
# return {"Hello": self.name}
|
|
|
|
# async def connect(self, ws: WebSocket):
|
|
# # 等待连接
|
|
# await ws.accept()
|
|
# # 存储ws连接对象
|
|
# self.active_connections.append(ws)
|
|
#
|
|
# def disconnect(self, ws: WebSocket):
|
|
# # 关闭时 移除ws对象
|
|
# self.active_connections.remove(ws)
|
|
#
|
|
# @staticmethod
|
|
# async def send_personal_message(message: str, ws: WebSocket):
|
|
# # 发送个人消息
|
|
# await ws.send_text(message)
|
|
#
|
|
# async def broadcast(self, message: str):
|
|
# # 广播消息
|
|
# for connection in self.active_connections:
|
|
# await connection.send_text(message)
|
|
#
|
|
#
|
|
# manager = ConnectionManager()
|
|
#
|
|
#
|
|
# @app.websocket("/ws/{user}")
|
|
# async def websocket_endpoint(websocket: WebSocket, user: str):
|
|
# await manager.connect(websocket)
|
|
#
|
|
# await manager.broadcast(f"用户{user}进入聊天室")
|
|
#
|
|
# try:
|
|
# while True:
|
|
# data = await websocket.receive_text()
|
|
# await manager.send_personal_message(f"你说了: {data}", websocket)
|
|
# await manager.broadcast(f"用户:{user} 说: {data}")
|
|
#
|
|
# except WebSocketDisconnect:
|
|
# manager.disconnect(websocket)
|
|
# await manager.broadcast(f"用户-{user}-离开")
|
|
|
|
|
|
class Web(Node):
|
|
topics = [ImageArgMsg]
|
|
|
|
def __init__(self, level=logging.INFO):
|
|
super(Web, self).__init__(level=level)
|
|
self.wss: list[WebSocket] = []
|
|
|
|
def custom_setup(self):
|
|
self.router = APIRouter()
|
|
self.router.add_api_route("/hello", self.hello, methods=["GET"])
|
|
self.router.add_websocket_route("/ws", self.websocket_endpoint)
|
|
|
|
def hello(self):
|
|
return {"Hello": 'asd'}
|
|
|
|
async def websocket_endpoint(self, websocket: WebSocket):
|
|
await websocket.accept()
|
|
self.wss.append(websocket)
|
|
t = ImageArgMsg('z', t_end=1000)
|
|
try:
|
|
while True:
|
|
data = await websocket.receive_json()
|
|
t.t_end = data['v']
|
|
self.send(t)
|
|
except WebSocketDisconnect:
|
|
self.wss.remove(websocket)
|
|
|
|
def boardcast(self, data: dict):
|
|
for ws in self.wss:
|
|
asyncio.run(ws.send_json(data))
|
|
|
|
def wst(self):
|
|
app = FastAPI()
|
|
app.include_router(self.router)
|
|
uvicorn.run(app=app, host="0.0.0.0", port=8010)
|
|
|
|
def loop(self):
|
|
Thread(target=self.wst).start()
|
|
while True:
|
|
msg = self.recv()
|
|
if isinstance(msg, ImageArgMsg):
|
|
if msg.sender != 'z':
|
|
self.boardcast(msg.__dict__)
|
|
logger.info(msg)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
Web()()
|