CameraMachine/python/sockets/socket_server.py

import json
import asyncio
from models import *
from .connect_manager import ConnectionManager
from .message_handler import *
from mcu.DeviceControl import DeviceControl,checkMcuConnection
from mcu.BlueToothMode import BlueToothMode
from mcu.capture.smart_shooter_class import SmartShooter
import time
from .socket_client import socket_manager
conn_manager = ConnectionManager()
active_connections = set()
device_ctrl = DeviceControl(websocket_manager=conn_manager)
blue_tooth = BlueToothMode(websocket_manager=conn_manager)
smart_shooter = SmartShooter(websocket_manager=conn_manager)
from utils.common import message_queue
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await conn_manager.connect(websocket)
    active_connections.add(websocket)
    smart_shooter.websocket = websocket
    try:
        # await socket_manager.connect()
        async def handler_messages():
            while True:
                try:
                    byteDats = await websocket.receive()
                    socket_type = byteDats.get("type")
                    if socket_type == "websocket.disconnect":
                        smart_shooter.stop_listen = True
                        device_ctrl.close_connect()
                        device_ctrl.mcu_exit = True
                        device_ctrl.clearMyInstance()
                        diviceList = blue_tooth.devices
                        if len(diviceList) == 0:
                            blue_tooth.bluetooth_exit = True
                            blue_tooth.clearMyInstance()
                            break
                        diviceAddress = list(diviceList.keys())[0]
                        if diviceAddress != "":
                            print(diviceList.get(diviceAddress))
                            diviceName = diviceList[diviceAddress]["name"]
                            blue_tooth.disconnect_device(diviceAddress, diviceName)
                        blue_tooth.bluetooth_exit = True
                        blue_tooth.clearMyInstance()
                        break
                    print("byteDats", byteDats)
                    await handlerSend(
                        conn_manager, json.dumps(byteDats), websocket, smart_shooter
                    )
                except Exception as e:
                    print(e)
                    break
        # async def checkConnMcu():
        #     await checkMcuConnection(device_ctrl)
        # async def connectBlueTooth():
        #     await blue_tooth.main_func()
        async def send_message():
            while True:
                try:
                    message = await message_queue.get()
                    await websocket.send_json(message)
                except Exception as e:
                    print(e)
                    break

        await asyncio.gather(handler_messages(), send_message())
    except WebSocketDisconnect:
        # socket_manager.close()
        print("Client disconnected")
    finally:
        active_connections.discard(websocket)
        # if websocket:
        #     await websocket.close()


@app.on_event("shutdown")
async def shutdown_event():
    print("Shutting down...")
    # socket_manager.close()
    # 清理操作
    for connection in list(active_connections):
        try:
            await connection.close()
        except Exception as e:
            print(f"Error closing connection: {e}")