Update pool.py

- Introduced new function `NodePool.get_player()` which gets the player after searching from all the nodes.
- Added the algorithim to the `NodePool.get_node()`
- Added config (its a dictionary for now, but it can be modified later on ), allows in customizing in toggling of settings
This commit is contained in:
Crussader 2021-10-26 21:06:31 +04:00 committed by GitHub
parent 292c9d18c7
commit 4d4587b179
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 42 additions and 7 deletions

View File

@ -6,7 +6,7 @@ import random
import re import re
import socket import socket
import time import time
from typing import Dict, Optional, TYPE_CHECKING from typing import Dict, Optional, TYPE_CHECKING, Union
from urllib.parse import quote from urllib.parse import quote
import aiohttp import aiohttp
@ -27,7 +27,12 @@ from .exceptions import (
TrackLoadError TrackLoadError
) )
from .objects import Playlist, Track from .objects import Playlist, Track
from .utils import ClientType, ExponentialBackoff, NodeStats from .utils import (
ClientType,
ExponentialBackoff,
NodeStats,
NodeAlgorithims,
)
if TYPE_CHECKING: if TYPE_CHECKING:
from .player import Player from .player import Player
@ -206,7 +211,7 @@ class Node:
if not self._available: if not self._available:
raise NodeNotAvailable( raise NodeNotAvailable(
f"The node '{self.identifier}' is not currently available.") f"The node '{self.identifier}' is not currently available.")
await self._websocket.send_str(json.dumps(data)) await self._websocket.send_str(json.dumps(data))
def get_player(self, guild_id: int): def get_player(self, guild_id: int):
@ -466,20 +471,50 @@ class NodePool:
def node_count(self): def node_count(self):
return len(self._nodes.values()) return len(self._nodes.values())
# @classmethod
# def get_node(cls, *, identifier: str = None) -> Node:
# """Fetches a node from the node pool using it's identifier.
# If no identifier is provided, it will choose a node at random.
# """
# available_nodes = {identifier: node for identifier, node in cls._nodes.items() if node._available}
# if not available_nodes:
# raise NoNodesAvailable('There are no nodes available.')
# if identifier is None:
# return random.choice(list(available_nodes.values()))
# return available_nodes.get(identifier, None)
@classmethod @classmethod
def get_node(cls, *, identifier: str = None) -> Node: def get_node(
cls, *,
identifier: str = None,
algorithim : NodeAlgorithims = NodeAlgorithims.base,
args=(), **kwargs
) -> Node:
"""Fetches a node from the node pool using it's identifier. """Fetches a node from the node pool using it's identifier.
If no identifier is provided, it will choose a node at random. If no identifier is provided, it will choose based on the algorithim given.
Default Algorithim returns a random Node.
""" """
available_nodes = {identifier: node for identifier, node in cls._nodes.items() if node._available} available_nodes = {identifier: node for identifier, node in cls._nodes.items() if node._available}
if not available_nodes: if not available_nodes:
raise NoNodesAvailable('There are no nodes available.') raise NoNodesAvailable('There are no nodes available.')
if identifier is None: if identifier is None and isinstance(algorithim, NodeAlgorithims):
return random.choice(list(available_nodes.values())) return algorithim(list(available_nodes.values()), *args, **kwargs).first
return available_nodes.get(identifier, None) return available_nodes.get(identifier, None)
@classmethod
def get_player(cls, guildId) -> Union[Player, None]:
"""Retruns the Exact Player object after Searching all the Nodes.
"""
for node in [node for node in cls._nodes.values() if node._available]:
if (player := node.players.get(guildId, None)):
return player
return None
@classmethod @classmethod
async def create_node( async def create_node(
cls, cls,