FastAPI 學習之路（五十）WebSockets（六）聊天室完善 FastAPI 學習之路（四十六）WebSockets（三）登錄後纔可以聊天

原創

2021-10-25 14:18

本次是WebSockets的第六次分享了，我們這次只是對於之前的功能做下優化，順便利用下之前的操作數據的接口，使用下數據庫的練習。

我們都知道，在聊天室裏面會有一個上線的概念。上線了要通知大家。然後下線了要告訴大家誰離開了。發小小大家要都能可以收到。那麼我們基於這個去完善下我們的功能。

首先，我們的登錄用戶還是用了之前分享中分享的數據庫相關呢，我們現在直接來調用，我們之前有從數據庫獲取用戶，然後校驗密碼，密碼成功我們才讓登錄成功。這個FastAPI 學習之路（四十六）WebSockets（三）登錄後纔可以聊天就已經分享了。

其次，我們去建立了鏈接，但是我們現在的鏈接都是沒有針對發消息做管理呢，我們去擴展下我們之前的發消息的類。

class ConnectionManager:
    def __init__(self):
        # 存放**的鏈接
        self.active_connections: List[Dict[str, WebSocket]] = []

    async def connect(self, user: str, ws: WebSocket):
        # 鏈接
        await ws.accept()
        self.active_connections.append({"user": user, "ws": ws})

    def disconnect(self, user: str, ws: WebSocket):
        # 關閉時 移除ws對象
        self.active_connections.remove({"user": user, "ws": ws})

    @staticmethod
    async def send_personal_message(message: str, ws: WebSocket):
        # 發送所有人消息
        await ws.send_text(message)

    async def send_other_message(self, message: dict, user: str):
        # 發送個人消息
        for connection in self.active_connections:
            if connection["user"] == user:
                await connection['ws'].send_json(message)

    async def broadcast(self, data: str):
        # 廣播消息
        for connection in self.active_connections:
            await connection['ws'].send_text(data)

增加了廣播消息，和增加了發所有消息和針對個人發消息。我們現在實現的是針對所有人和廣播。那麼我們看下，我們的管理聊天鏈接的已經完善。

接下來我們就是去實現如何接受消息。

@app.websocket("/items/ws/{user}")
async def websocket_endpoint(
        websocket: WebSocket,
        user:str,
        cookie_or_token: str = Depends(get_cookie_or_token),

):
    await manager.connect(user, websocket)
    await manager.broadcast(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 as e:
        manager.disconnect(user, websocket)
        await manager.broadcast("用戶{}離開".format(user))

其實很簡單，當你鏈接的時候，我們就廣告一下，誰誰進入了聊天室。然後接着就是等消息。那麼我們看下我們之前寫的前端需要怎麼改造呢

<!DOCTYPE html>
<html>
<head>
    <title>Chat</title>
</head>
<body>
<h1>WebSocket 聊天</h1>
<form action="" onsubmit="sendMessage(event)">
    <input type="text" id="messageText" autocomplete="off"/>
    <button>Send</button>
</form>
<button onclick="logout()">退出</button>
<ul id='messages'>
</ul>
<script>
    var  token=window.localStorage.getItem("token")
    if (token==null ){
        window.location.href="/login"
    }
    var ws = new WebSocket("ws://localhost:8000/items/ws/"+token+"?token="+token);

    ws.onmessage = function (event) {

        var messages = document.getElementById('messages')

        var message = document.createElement('li')

        var content = document.createTextNode(event.data)

        message.appendChild(content)

        messages.appendChild(message)

    };

    function sendMessage(event) {

        var input = document.getElementById("messageText")

        ws.send(input.value)

        input.value = ''

        event.preventDefault()

    }
    function logout() {
        window.localStorage.removeItem("token")
        window.location.href='/login'
    }
</script>

</body>

</html>

其實我們就是在連接的時候增加了對應人的連接，那麼我們看下，我們應該測試下效果。