fixed all warnings, and added location algorithm

pomice/enums.py

@@ -129,6 +129,7 @@ class NodeAlgorithm(Enum):
     # We don't have to define anything special for these, since these just serve as flags
     by_ping = "BY_PING"
     by_players = "BY_PLAYERS"
+    by_location = "BY_LOCATION"
 
     def __str__(self) -> str:
         return self.value

pomice/pool.py

@@ -18,9 +18,8 @@ from urllib.parse import quote
 
 import aiohttp
 import orjson as json
-from discord import Client
+from discord import Client, VoiceChannel
 from discord.ext import commands
-from discord.utils import MISSING
 
 try:
     from websockets.legacy import client  # websockets >= 10.0
@@ -28,13 +27,17 @@ except ImportError:
     import websockets.client as client  # websockets < 10.0  # type: ignore
 
 from websockets import exceptions
-from websockets import typing as wstype
 
 from . import __version__
 from . import applemusic
 from . import spotify
-from .enums import *
+from .enums import (
-from .enums import LogLevel
+    TrackType,
+    NodeAlgorithm,
+    PlaylistType,
+    SearchType,
+    URLRegex
+)
 from .exceptions import InvalidSpotifyClientAuthorization
 from .exceptions import LavalinkVersionIncompatible
 from .exceptions import NodeConnectionFailure
@@ -78,6 +81,7 @@ class Node:
         "_pool",
         "_password",
         "_identifier",
+        "_location",
         "_heartbeat",
         "_resume_key",
         "_resume_timeout",
@@ -114,6 +118,7 @@ class Node:
         port: int,
         password: str,
         identifier: str,
+        location: str = "us-east",
         secure: bool = False,
         heartbeat: int = 120,
         resume_key: Optional[str] = None,
@@ -141,6 +146,8 @@ class Node:
         self._secure: bool = secure
         self._fallback: bool = fallback
         
+        self._location = location
+
         self._websocket_uri: str = f"{'wss' if self._secure else 'ws'}://{self._host}:{self._port}"
         self._rest_uri: str = f"{'https' if self._secure else 'http'}://{self._host}:{self._port}"
 
@@ -212,6 +219,14 @@ class Node:
         """Property which returns the discord.py client linked to this node"""
         return self._bot
         
+    @property
+    def location(self) -> str:
+        """
+        Property which returns the default region unless set specifically
+        """
+        
+        return self._location
+
     @property
     def player_count(self) -> int:
         """Property which returns how many players are connected to this node"""
@@ -366,7 +381,7 @@ class Node:
             self._session_id = data["sessionId"]
             await self._configure_resuming()
 
-        if not "guildId" in data:
+        if "guildId" not in data:
             return
 
         player: Optional[Player] = self._players.get(int(data["guildId"]))
@@ -969,7 +984,7 @@ class NodePool:
         return len(self._nodes.values())
 
     @classmethod
-    def get_best_node(cls, *, algorithm: NodeAlgorithm) -> Node:
+    def get_best_node(cls, *, algorithm: NodeAlgorithm, channel: Optional[VoiceChannel] = None) -> Node:
         """Fetches the best node based on an NodeAlgorithm.
         This option is preferred if you want to choose the best node
         from a multi-node setup using either the node's latency
@@ -983,6 +998,7 @@ class NodePool:
         based on how players it has. This method will return a node with
         the least amount of players
         """
+        
         available_nodes: List[Node] = [node for node in cls._nodes.values() if node._available]
 
         if not available_nodes:
@@ -996,6 +1012,20 @@ class NodePool:
             tested_nodes = {node: len(node.players.keys()) for node in available_nodes}
             return min(tested_nodes, key=tested_nodes.get)  # type: ignore
         
+        elif algorithm == NodeAlgorithm.by_location and isinstance(channel, VoiceChannel):
+            tested_nodes = {} 
+            chosen_region = channel.rtc_region
+            
+            if not chosen_region:
+                return cls.get_best_node(algorithm=NodeAlgorithm.by_ping)
+                
+            if (node := next(
+                (node for node in available_nodes if node.location == chosen_region),
+                None
+            )):
+                return node
+            
+            return random.choice(available_nodes)
         else:
             raise ValueError(
                 "The algorithm provided is not a valid NodeAlgorithm.",
@@ -1027,6 +1057,7 @@ class NodePool:
         port: int,
         password: str,
         identifier: str,
+        location: str = "us-east",
         secure: bool = False,
         heartbeat: int = 120,
         resume_key: Optional[str] = None,
@@ -1054,6 +1085,7 @@ class NodePool:
             port=port,
             password=password,
             identifier=identifier,
+            location=location,
             secure=secure,
             heartbeat=heartbeat,
             resume_key=resume_key,

pyproject.toml

@@ -1,3 +1,16 @@
+[project]
+name = "pomice"
+version = "2.10.0"
+description = "The modern Lavalink wrapper designed for Discord.py"
+readme = "README.md"
+requires-python = ">=3.8"
+dependencies = [
+    "aiohttp>=3.7.4,<4",
+    "discord>=2.3.2",
+    "orjson",
+    "websockets",
+]
+
 [build-system]
 requires = [
     "setuptools>=42",
@@ -17,3 +30,6 @@ no_implicit_optional = true
 check_untyped_defs = true
 warn_unused_ignores = true
 show_error_codes = true
+
+[tool.pyright]
+typeCheckingMode = "standard"