god pydantic is nice

nightwatch/server/__init__.py

@@ -1,7 +1,6 @@
 # Copyright (c) 2024 iiPython
 
 # Modules
-from typing import Type
 import orjson
 from websockets import WebSocketCommonProtocol
 from websockets.exceptions import ConnectionClosedError
@@ -18,11 +17,9 @@ def __init__(self) -> None:
         self.message_buffer = []
 
     def add_client(self, client: WebSocketCommonProtocol) -> None:
-        print("Client registered")
         self.clients.add(client)
 
     def remove_client(self, client: WebSocketCommonProtocol) -> None:
-        print("Client removed")
         self.clients.remove(client)
 
 state = NightwatchStateManager()
@@ -39,20 +36,8 @@ async def connection(websocket: WebSocketCommonProtocol) -> None:
                 await client.send("error", text = "Specified command type does not exist or is missing.")
                 continue
 
-            data = message.get("data") or {}
-            missing_parameters = registry.types[message["type"]].keys() - data.keys()
-            if missing_parameters:
-                await client.send("error", text = f"The following parameters are missing: {', '.join(missing_parameters)}.")
-                continue
-
-            try:
-                await registry.commands[message["type"]](state, client, *[
-                    v(data[k])
-                    for k, v in registry.types[message["type"]].items()
-                ])
-
-            except (TypeError, ValueError):
-                await client.send("error", text = "Failed to convert your arguments to the appropriate datatype.")
+            command, payload_type = registry.commands[message["type"]]
+            await command(state, client, payload_type(**(message.get("data") or {})))
 
     except orjson.JSONDecodeError:
         log.warn("ws", "Failed to decode JSON from client.")

nightwatch/server/utils/commands.py

@@ -3,24 +3,23 @@
 # Modules
 from typing import Callable
 
+from . import models
+from .websocket import NightwatchClient
+
 # Handle command registration
 class CommandRegistry():
     def __init__(self) -> None:
-        self.types, self.commands = {}, {}
+        self.commands = {}
 
     def command(self, name: str) -> Callable:
         def callback(function: Callable) -> None:
-            self.types[name] = {
-                k: v for k, v in function.__annotations__.items()
-                if k not in ["state", "client", "return"]
-            }
-            self.commands[name] = function
+            self.commands[name] = (function, function.__annotations__["data"])
 
         return callback
 
 registry = CommandRegistry()
 
 # Setup commands
 @registry.command("identify")
-async def command_identify(state, client, name: str, color: int) -> None:
-    print(state, client, name, color)
+async def command_identify(state, client: NightwatchClient, data: models.IdentifyModel) -> None:
+    print(data)

nightwatch/server/utils/models.py

@@ -0,0 +1,11 @@
+# Copyright (c) 2024 iiPython
+
+# Modules
+from typing import Annotated
+from pydantic import BaseModel, StringConstraints
+from pydantic_extra_types.color import Color
+
+# Models
+class IdentifyModel(BaseModel):
+    name: Annotated[str, StringConstraints(min_length = 3, max_length = 32)]
+    color: Color

pyproject.toml

@@ -37,6 +37,8 @@ dependencies = [
 serve = [
     "fastapi",
     "pymongo",
+    "pydantic",
+    "pydantic-extra-types",
     "starlette",
     "argon2-cffi",
     "redis[hiredis]",