106 lines
3.0 KiB
Python
106 lines
3.0 KiB
Python
import asyncio
|
|
import logging
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
|
|
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
|
|
from fastapi.responses import FileResponse, JSONResponse
|
|
|
|
from ros_node import get_node, start_ros, stop_ros
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
STATIC_DIR = Path(__file__).parent / 'static'
|
|
|
|
|
|
class ConnectionManager:
|
|
def __init__(self) -> None:
|
|
self.connections: list[WebSocket] = []
|
|
self._loop: asyncio.AbstractEventLoop | None = None
|
|
|
|
def set_loop(self, loop: asyncio.AbstractEventLoop) -> None:
|
|
self._loop = loop
|
|
|
|
async def connect(self, ws: WebSocket) -> None:
|
|
await ws.accept()
|
|
self.connections.append(ws)
|
|
|
|
def disconnect(self, ws: WebSocket) -> None:
|
|
if ws in self.connections:
|
|
self.connections.remove(ws)
|
|
|
|
def broadcast_from_thread(self, data: dict) -> None:
|
|
if self._loop and self._loop.is_running():
|
|
asyncio.run_coroutine_threadsafe(self._broadcast(data), self._loop)
|
|
|
|
async def _broadcast(self, data: dict) -> None:
|
|
dead = []
|
|
for ws in self.connections:
|
|
try:
|
|
await ws.send_json(data)
|
|
except Exception:
|
|
dead.append(ws)
|
|
for ws in dead:
|
|
self.disconnect(ws)
|
|
|
|
|
|
manager = ConnectionManager()
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
manager.set_loop(asyncio.get_event_loop())
|
|
node = start_ros()
|
|
node.on_joint_states_callbacks.append(manager.broadcast_from_thread)
|
|
node.on_ultrasonic_callbacks.append(manager.broadcast_from_thread)
|
|
yield
|
|
stop_ros()
|
|
|
|
|
|
app = FastAPI(lifespan=lifespan)
|
|
|
|
|
|
@app.websocket('/ws')
|
|
async def websocket_endpoint(ws: WebSocket):
|
|
await manager.connect(ws)
|
|
node = get_node()
|
|
try:
|
|
while True:
|
|
data = await ws.receive_json()
|
|
msg_type = data.get('type')
|
|
if msg_type == 'cmd_vel':
|
|
node.publish_cmd_vel(data.get('linear_x', 0.0), data.get('angular_z', 0.0))
|
|
elif msg_type == 'joint_command':
|
|
node.publish_joint_command(data.get('pan', 0.0), data.get('tilt', 0.0))
|
|
except WebSocketDisconnect:
|
|
pass
|
|
except Exception as e:
|
|
logger.error('WebSocket error: %s', e)
|
|
finally:
|
|
manager.disconnect(ws)
|
|
try:
|
|
node.publish_cmd_vel(0.0, 0.0)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
@app.get('/{full_path:path}')
|
|
async def serve_spa(full_path: str):
|
|
index = STATIC_DIR / 'index.html'
|
|
if not index.exists():
|
|
return JSONResponse(
|
|
{'detail': 'Frontend not built. In dev mode open http://localhost:5173'},
|
|
status_code=404,
|
|
)
|
|
candidate = (STATIC_DIR / full_path).resolve()
|
|
static_root = STATIC_DIR.resolve()
|
|
if str(candidate).startswith(str(static_root)) and full_path and candidate.is_file():
|
|
return FileResponse(candidate)
|
|
return FileResponse(index)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
import uvicorn
|
|
uvicorn.run(app, host='0.0.0.0', port=8080)
|