♻️ Cleaned socket_server up by using classes instead of dicts

classquiz/socket_server/__init__.py

@@ -4,8 +4,10 @@ import os
 import aiohttp
 import socketio
 
+from typing import Any
 from classquiz.config import redis, settings
 from classquiz.db.models import PlayGame
+from pydantic import BaseModel, ValidationError
 
 sio = socketio.AsyncServer(async_mode="asgi", cors_allowed_origins=[])
 settings = settings()
@@ -22,16 +24,44 @@ async def generate_final_results(game_data: PlayGame, game_pin: str) -> dict:
     return results
 
 
+class _JoinGameData(BaseModel):
+    username: str
+    game_pin: str
+    captcha: str | None
+
+
+class _GameSessionPlayer(BaseModel):
+    username: str
+    sid: str
+
+
+class _GameSession(BaseModel):
+    admin: str
+    game_id: str
+    players: list[_GameSessionPlayer | None]
+    answers: list[Any]
+
+
 @sio.event
 async def join_game(sid: str, data: dict):
+    redis_res = await redis.get(f"game:{data['game_pin']}")
+    if redis_res is None:
+        await sio.emit("game_not_found", room=sid)
+        return
+    try:
+        data = _JoinGameData(**data)
+    except ValidationError as e:
+        await sio.emit("error", room=sid)
+        print(e)
+        return
+    # +++ START checking captcha +++
     async with aiohttp.ClientSession() as session:
-        redis_res = await redis.get(f"game:{data['game_pin']}")
         try:
             if json.loads(redis_res)["captcha_enabled"]:
                 try:
                     async with session.post(
                         "https://hcaptcha.com/siteverify",
-                        data={"response": data["captcha"], "secret": settings.hcaptcha_key},
+                        data={"response": data.captcha, "secret": settings.hcaptcha_key},
                     ) as resp:
                         resp_data = await resp.json()
                         if not resp_data["success"]:
@@ -43,37 +73,28 @@ async def join_game(sid: str, data: dict):
                     return
         except TypeError:
             pass
-    if redis_res is None:
+    # --- END checking captcha ---
-        await sio.emit("game_not_found", room=sid)
+    session = {
-    else:
+        "game_pin": data.game_pin,
-        session = {
+        "username": data.username,
-            "game_pin": data["game_pin"],
+        "admin": False,
-            "username": data["username"],
+    }
-            "admin": False,
+    await sio.save_session(sid, session)
-        }
+    await sio.emit("joined_game", redis_res, room=sid)
-        await sio.save_session(sid, session)
+    redis_res = await redis.get(f"game_session:{data.game_pin}")
-        await sio.emit("joined_game", redis_res, room=sid)
+    redis_res = _GameSession.parse_raw(redis_res)
-        redis_res = await redis.get(f"game_session:{data['game_pin']}")
+    redis_res.players.append(_GameSessionPlayer(username=data.username, sid=sid))
-        redis_res = json.loads(redis_res)
+    await redis.set(
-        redis_res["players"].append({"username": data["username"], "sid": sid})
+        f"game_session:{data.game_pin}",
-        await redis.set(
+        _GameSession(admin=redis_res.admin, game_id=redis_res.game_id, players=redis_res.players, answers=[]).json(),
-            f"game_session:{data['game_pin']}",
+        ex=18000,
-            json.dumps(
+    )
-                {
+    await sio.emit(
-                    "admin": redis_res["admin"],
+        "player_joined",
-                    "game_id": redis_res["game_id"],
+        {"username": data.username, "sid": sid},
-                    "players": redis_res["players"],
+        room=redis_res.admin,
-                    "answers": [],
+    )
-                }
+    sio.enter_room(sid, data.game_pin)
-            ),
-            ex=18000,
-        )
-        await sio.emit(
-            "player_joined",
-            {"username": data["username"], "sid": sid},
-            room=redis_res["admin"],
-        )
-        sio.enter_room(sid, data["game_pin"])
 
 
 @sio.event
@@ -83,26 +104,37 @@ async def start_game(sid: str, _data: dict):
         await sio.emit("start_game", room=session["game_pin"])
 
 
+class _RegisterAsAdminData(BaseModel):
+    game_pin: str
+    game_id: str
+
+
 @sio.event
 async def register_as_admin(sid: str, data: dict):
-    game_pin = data["game_pin"]
+    try:
-    game_id = data["game_id"]
+        data = _RegisterAsAdminData(**data)
+    except ValidationError as e:
+        await sio.emit("error", room=sid)
+        print(e)
+        return
+    game_pin = data.game_pin
+    game_id = data.game_id
     if (await redis.get(f"game_session:{game_pin}")) is None:
         await redis.set(
             f"game_session:{game_pin}",
-            json.dumps({"admin": sid, "game_id": game_id, "players": []}),
+            _GameSession(admin=sid, game_id=game_id, answers=[], players=[]).json(),
             ex=18000,
         )
 
         await sio.emit(
             "registered_as_admin",
-            {"game_id": game_id, "game": await redis.get(f"game:{data['game_pin']}")},
+            {"game_id": game_id, "game": await redis.get(f"game:{game_pin}")},
             room=sid,
         )
         async with sio.session(sid) as session:
-            session["game_pin"] = data["game_pin"]
+            session["game_pin"] = game_pin
             session["admin"] = True
-        sio.enter_room(sid, data["game_pin"])
+        sio.enter_room(sid, game_pin)
     else:
         await sio.emit("already_registered_as_admin", room=sid)
 
@@ -117,48 +149,59 @@ async def get_question_results(sid: str, data: dict):
 
 
 @sio.event
-async def set_question_number(sid, data):
+async def set_question_number(sid, data: str):
+    # data is just a number (as a str) of the question
     session = await sio.get_session(sid)
     if session["admin"]:
         game_pin = session["game_pin"]
         await sio.emit("set_question_number", data, room=game_pin)
 
 
+class _SubmitAnswerData(BaseModel):
+    question_index: int
+    answer: str
+
+
+class _AnswerData(BaseModel):
+    username: str
+    answer: str
+    right: bool
+
+
+class _AnswerDataList(BaseModel):
+    # Just a method to make a top-level list
+    __root__: list[_AnswerData]
+
+
 @sio.event
 async def submit_answer(sid: str, data: dict):
+    try:
+        data = _SubmitAnswerData(**data)
+    except ValidationError as e:
+        await sio.emit("error", room=sid)
+        print(e)
+        return
     session = await sio.get_session(sid)
-    game_data = PlayGame(**json.loads(await redis.get(f"game:{session['game_pin']}")))
+    game_data = PlayGame.parse_raw(await redis.get(f"game:{session['game_pin']}"))
     answer_right = False
-    for answer in game_data.questions[int(data["question_index"])].answers:
+    for answer in game_data.questions[int(data.question_index)].answers:
-        if answer.answer == data["answer"] and answer.right:
+        if answer.answer == data.answer and answer.right:
             answer_right = True
             break
-    answers = await redis.get(f"game_session:{session['game_pin']}:{data['question_index']}")
+    answers = await redis.get(f"game_session:{session['game_pin']}:{data.question_index}")
     if answers is None:
         await redis.set(
-            f"game_session:{session['game_pin']}:{data['question_index']}",
+            f"game_session:{session['game_pin']}:{data.question_index}",
-            json.dumps(
+            _AnswerDataList(
-                [
+                __root__=[_AnswerData(username=session["username"], answer=data.answer, right=answer_right)]
-                    {
+            ).json(),
-                        "username": session["username"],
-                        "answer": data["answer"],
-                        "right": answer_right,
-                    }
-                ]
-            ),
             ex=18000,
         )
     else:
-        answers = json.loads(answers)
+        answers = _AnswerDataList.parse_raw(answers)
-        answers.append(
+        answers.__root__.append(_AnswerData(username=session["username"], answer=data.answer, right=answer_right))
-            {
-                "username": session["username"],
-                "answer": data["answer"],
-                "right": answer_right,
-            }
-        )
         await redis.set(
-            f"game_session:{session['game_pin']}:{data['question_index']}",
+            f"game_session:{session['game_pin']}:{data.question_index}",
             json.dumps(answers),
             ex=18000,
         )