Refactoring, added docstrings

dragonion_server/modules/server/handlers/managers/connection.py

@@ -15,12 +15,22 @@ class Connection(object):
     public_key: str
 
     async def send_webmessage(self, obj: webmessages_union):
+        """
+        Sends WebMessage object to this connection
+        :param obj: Should be some type of WebMessage
+        :return: 
+        """
         await self.ws.send_text(obj.to_json())
 
     async def send_error(
             self,
             error_message: webmessage_error_message_literal
     ):
+        """
+        Sends error with specified messages
+        :param error_message: See webmessage_error_message_literal for available
+        :return: 
+        """
         await self.send_webmessage(
             WebErrorMessage(
                 error_message=error_message
@@ -28,6 +38,10 @@ class Connection(object):
         )
 
     async def send_connect(self):
+        """
+        When new user is connected, send info about user
+        :return: 
+        """
         await self.send_webmessage(
             WebUserMessage(
                 type="connect",

dragonion_server/modules/server/handlers/managers/room.py

@@ -18,6 +18,12 @@ class Room(object):
     connections: Dict[str, Connection] = {}
 
     async def accept_connection(self, ws: WebSocket) -> Connection:
+        """
+        Accepts connection, checks username availability and adds it to dict of 
+        connections 
+        :param ws: Websocket of connection 
+        :return: 
+        """
         print('Incoming connection')
         await ws.accept()
         connection = Connection(
@@ -36,11 +42,22 @@ class Room(object):
         return connection
 
     async def broadcast_webmessage(self, obj: webmessages_union):
+        """
+        Broadcasts WebMessages to all connections in room
+        :param obj:  
+        :return: 
+        """
         for connection in self.connections.values():
             print(f'Sending to {connection.username}: {obj}')
             await connection.send_webmessage(obj)
 
     async def broadcast_message(self, from_username: str, message: str):
+        """
+        Broadcasts message to every user in room
+        :param from_username: User that sent message
+        :param message: content
+        :return: 
+        """
         await self.broadcast_webmessage(
             WebMessageMessage(
                 username=from_username,
@@ -49,6 +66,11 @@ class Room(object):
         )
 
     async def broadcast_notification(self, message: str):
+        """
+        Broadcasts notification from server
+        :param message: Content
+        :return: 
+        """
         await self.broadcast_webmessage(
             WebNotificationMessage(
                 message=message
@@ -59,6 +81,11 @@ class Room(object):
             self,
             error_message: webmessage_error_message_literal
     ):
+        """
+        Broadcasts server error
+        :param error_message: See webmessage_error_message_literal
+        :return: 
+        """
         await self.broadcast_webmessage(
             WebErrorMessage(
                 error_message=error_message
@@ -66,6 +93,11 @@ class Room(object):
         )
 
     async def broadcast_user_disconnected(self, username: str):
+        """
+        Broadcasts that user is disconnected
+        :param username: Username of user that disconnected
+        :return: 
+        """
         await self.broadcast_webmessage(
             WebUserMessage(
                 type="disconnect",
@@ -74,11 +106,24 @@ class Room(object):
         )
 
     async def get_connection_by(self, attribute: str, value: str) -> Connection | None:
+        """
+        Search for connection by attribute and value in it
+        :param attribute: 
+        :param value: 
+        :return: 
+        """
         for connection in self.connections.values():
             if getattr(connection, attribute) == value:
                 return connection
 
     async def disconnect(self, connection: Connection, close_reason: str | None = None):
+        """
+        Disconnects by connection object. 
+        :param connection: Object of connection. 
+                           It can be obtained using get_connection_by
+        :param close_reason: Reason if exists
+        :return: 
+        """
         if connection not in self.connections.values():
             return
 

dragonion_server/modules/server/handlers/managers/service.py

@@ -29,6 +29,11 @@ class Service(object):
             await room.broadcast_error(error_message)
 
     async def get_room_by_connection(self, connection: Connection) -> Room:
+        """
+        Searches for room by valid connection object in it
+        :param connection: Connection in unknown room to search
+        :return: 
+        """
         for room in self.rooms.values():
             if connection in room.connections.values():
                 return room
@@ -36,11 +41,23 @@ class Service(object):
     async def get_connection_by_attribute(
             self, attribute: str, value: str
     ) -> Connection:
+        """
+        Gets connection in some room by attribute and value in it
+        :param attribute: 
+        :param value: 
+        :return: 
+        """
         for room in self.rooms.values():
             if connection := await room.get_connection_by(attribute, value):
                 return connection
 
     async def close_room(self, room_name: str, reason: str = 'Unknown reason'):
+        """
+        Closes all connections in room
+        :param room_name: Close name
+        :param reason: Reason to close room, default is Unknown reason
+        :return: 
+        """
         room = self.rooms.get(room_name)
         if room is None:
             return

dragonion_server/modules/server/handlers/websocket_server.py

@@ -10,6 +10,12 @@ service = Service()
 
 
 async def serve_websocket(websocket: WebSocket, room_name: str):
+    """
+    Serves websocket
+    :param websocket: Ws to serve 
+    :param room_name: Room name to connect ws to
+    :return: 
+    """
     print(f'Connection opened room {room_name}')
     room = await service.get_room(room_name)
     connection = await room.accept_connection(websocket)