Compare commits

..

No commits in common. "main" and "2.5.1" have entirely different histories.
main ... 2.5.1

18 changed files with 338 additions and 941 deletions

1
.gitignore vendored
View File

@ -10,7 +10,6 @@ build/
Pipfile.lock
.mypy_cache/
.vscode/
.idea/
.venv/
*.code-workspace
*.ini

View File

@ -2,7 +2,7 @@
# See https://pre-commit.com/hooks.html for more hooks
repos:
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v4.5.0
rev: v4.4.0
hooks:
- id: check-ast
- id: check-builtin-literals
@ -11,23 +11,31 @@ repos:
- id: requirements-txt-fixer
- id: trailing-whitespace
- repo: https://github.com/psf/black
rev: 23.10.1
rev: 23.3.0
hooks:
- id: black
language_version: python3.13
language_version: python3.10
- repo: https://github.com/asottile/blacken-docs
rev: 1.13.0
hooks:
- id: blacken-docs
- repo: https://github.com/asottile/pyupgrade
rev: v3.15.0
rev: v3.3.2
hooks:
- id: pyupgrade
args: [--py37-plus, --keep-runtime-typing]
- repo: https://github.com/asottile/reorder-python-imports
rev: v3.12.0
- repo: https://github.com/asottile/reorder_python_imports
rev: v3.9.0
hooks:
- id: reorder-python-imports
- repo: https://github.com/asottile/add-trailing-comma
rev: v3.1.0
rev: v2.4.0
hooks:
- id: add-trailing-comma
- repo: https://github.com/hadialqattan/pycln
rev: v2.1.3
hooks:
- id: pycln
default_language_version:
python: python3.13
python: python3.10

View File

@ -6,7 +6,6 @@ name = "pypi"
[packages]
orjson = "*"
"discord.py" = {extras = ["voice"], version = "*"}
websockets = "*"
[dev-packages]
mypy = "*"

View File

@ -13,9 +13,11 @@ The classes listed here are as they appear in Pomice. When you use them within y
the way you use them will be different. Here's an example on how you would use the `TrackStartEvent` within an event listener in a cog:
```py
@commands.Cog.listener
async def on_pomice_track_start(self, player: Player, track: Track):
...
```
## Event definitions

View File

@ -66,10 +66,13 @@ After you have initialized your function, we need to fill in the proper paramete
- Set this value to `True` if you want Pomice to automatically switch all players to another available node if one disconnects.
You must have two or more nodes to be able to do this.
* - `logger`
- `Optional[logging.Logger]`
- If you would like to receive logging information from Pomice, set this to your logger class
* - `log_level`
- `LogLevel`
- The logging level for the node. The default logging level is `LogLevel.INFO`.
* - `log_handler`
- `Optional[logging.Handler]`
- The logging handler for the node. Set to `None` to default to the built-in logging handler.
:::
@ -90,13 +93,13 @@ await NodePool.create_node(
spotify_client_secret="<your spotify client secret here>"
apple_music=<True/False>,
fallback=<True/False>,
logger=<your logger here>
log_level=<optional LogLevel here>
)
```
:::{important}
For features like Spotify and Apple Music, you are **not required** to fill in anything for them if you do not want to use them. If you do end up queuing a Spotify or Apple Music track, it is **up to you** on how you decide to handle it, whether it be through your own methods or a Lavalink plugin.
For features like Spotify and Apple Music, you are **not required** to fill in anything for them if you do not want to use them. If you do end up queuing a Spotify or Apple Music track anyway, they will **not work** because these options are not enabled.
:::

View File

@ -10,71 +10,63 @@ import re
from discord.ext import commands
URL_REG = re.compile(r"https?://(?:www\.)?.+")
URL_REG = re.compile(r'https?://(?:www\.)?.+')
class MyBot(commands.Bot):
def __init__(self) -> None:
super().__init__(
command_prefix="!",
activity=discord.Activity(
type=discord.ActivityType.listening, name="to music!"
),
)
self.add_cog(Music(self))
def __init__(self) -> None:
super().__init__(command_prefix='!', activity=discord.Activity(type=discord.ActivityType.listening, name='to music!'))
async def on_ready(self) -> None:
print("I'm online!")
await self.cogs["Music"].start_nodes()
self.add_cog(Music(self))
async def on_ready(self) -> None:
print("I'm online!")
await self.cogs["Music"].start_nodes()
class Music(commands.Cog):
def __init__(self, bot) -> None:
self.bot = bot
self.pomice = pomice.NodePool()
def __init__(self, bot) -> None:
self.bot = bot
async def start_nodes(self):
await self.pomice.create_node(
bot=self.bot,
host="127.0.0.1",
port="3030",
password="youshallnotpass",
identifier="MAIN",
)
print(f"Node is ready!")
self.pomice = pomice.NodePool()
@commands.command(name="join", aliases=["connect"])
async def join(
self, ctx: commands.Context, *, channel: discord.TextChannel = None
) -> None:
if not channel:
channel = getattr(ctx.author.voice, "channel", None)
async def start_nodes(self):
await self.pomice.create_node(bot=self.bot, host='127.0.0.1', port='3030',
password='youshallnotpass', identifier='MAIN')
print(f"Node is ready!")
@commands.command(name='join', aliases=['connect'])
async def join(self, ctx: commands.Context, *, channel: discord.TextChannel = None) -> None:
if not channel:
channel = getattr(ctx.author.voice, 'channel', None)
if not channel:
raise commands.CheckFailure(
"You must be in a voice channel to use this command"
"without specifying the channel argument."
)
raise commands.CheckFailure('You must be in a voice channel to use this command'
'without specifying the channel argument.')
await ctx.author.voice.channel.connect(cls=pomice.Player)
await ctx.send(f"Joined the voice channel `{channel}`")
@commands.command(name="play")
async def play(self, ctx, *, search: str) -> None:
if not ctx.voice_client:
await ctx.author.voice.channel.connect(cls=pomice.Player)
await ctx.send(f'Joined the voice channel `{channel}`')
@commands.command(name='play')
async def play(self, ctx, *, search: str) -> None:
if not ctx.voice_client:
await ctx.invoke(self.join)
player = ctx.voice_client
player = ctx.voice_client
results = await player.get_tracks(query=f"{search}")
results = await player.get_tracks(query=f'{search}')
if not results:
raise commands.CommandError("No results were found for that search term.")
if not results:
raise commands.CommandError('No results were found for that search term.')
if isinstance(results, pomice.Playlist):
if isinstance(results, pomice.Playlist):
await player.play(track=results.tracks[0])
else:
else:
await player.play(track=results[0])

View File

@ -3,4 +3,3 @@ discord.py[voice]
furo
myst_parser
orjson
websockets

View File

@ -3,7 +3,7 @@ Pomice
~~~~~~
The modern Lavalink wrapper designed for discord.py.
Copyright (c) 2024, cloudwithax
Copyright (c) 2023, cloudwithax
Licensed under GPL-3.0
"""
@ -20,7 +20,7 @@ if not discord.version_info.major >= 2:
"using 'pip install discord.py'",
)
__version__ = "2.10.0"
__version__ = "2.5.1"
__title__ = "pomice"
__author__ = "cloudwithax"
__license__ = "GPL-3.0"

View File

@ -1,14 +1,11 @@
from __future__ import annotations
import asyncio
import base64
import logging
import re
from datetime import datetime
from typing import AsyncGenerator
from typing import Dict
from typing import List
from typing import Optional
from typing import Union
import aiohttp
@ -20,10 +17,10 @@ from .objects import *
__all__ = ("Client",)
AM_URL_REGEX = re.compile(
r"https?://music\.apple\.com/(?P<country>[a-zA-Z]{2})/(?P<type>album|playlist|song|artist)/(?P<name>.+?)/(?P<id>[^/?]+?)(?:/)?(?:\?.*)?$",
r"https?://music.apple.com/(?P<country>[a-zA-Z]{2})/(?P<type>album|playlist|song|artist)/(?P<name>.+)/(?P<id>[^?]+)",
)
AM_SINGLE_IN_ALBUM_REGEX = re.compile(
r"https?://music\.apple\.com/(?P<country>[a-zA-Z]{2})/(?P<type>album|playlist|song|artist)/(?P<name>.+)/(?P<id>[^/?]+)(\?i=)(?P<id2>[^&]+)(?:&.*)?$",
r"https?://music.apple.com/(?P<country>[a-zA-Z]{2})/(?P<type>album|playlist|song|artist)/(?P<name>.+)/(?P<id>.+)(\?i=)(?P<id2>.+)",
)
AM_SCRIPT_REGEX = re.compile(r'<script.*?src="(/assets/index-.*?)"')
@ -38,14 +35,12 @@ class Client:
and translating it to a valid Lavalink track. No client auth is required here.
"""
def __init__(self, *, playlist_concurrency: int = 6) -> None:
def __init__(self) -> None:
self.expiry: datetime = datetime(1970, 1, 1)
self.token: str = ""
self.headers: Dict[str, str] = {}
self.session: aiohttp.ClientSession = None # type: ignore
self._log = logging.getLogger(__name__)
# Concurrency knob for parallel playlist page retrieval
self._playlist_concurrency = max(1, playlist_concurrency)
async def _set_session(self, session: aiohttp.ClientSession) -> None:
self.session = session
@ -101,8 +96,7 @@ class Client:
).decode()
token_data = json.loads(token_json)
self.expiry = datetime.fromtimestamp(token_data["exp"])
if self._log:
self._log.debug(f"Fetched Apple Music bearer token successfully")
self._log.debug(f"Fetched Apple Music bearer token successfully")
async def search(self, query: str) -> Union[Album, Playlist, Song, Artist]:
if not self.token or datetime.utcnow() > self.expiry:
@ -136,10 +130,9 @@ class Client:
)
data: dict = await resp.json(loads=json.loads)
if self._log:
self._log.debug(
f"Made request to Apple Music API with status {resp.status} and response {data}",
)
self._log.debug(
f"Made request to Apple Music API with status {resp.status} and response {data}",
)
data = data["data"][0]
@ -172,127 +165,25 @@ class Client:
"This playlist is empty and therefore cannot be queued.",
)
# Apple Music uses cursor pagination with 'next'. We'll fetch subsequent pages
# concurrently by first collecting cursors in rolling waves.
next_cursor = track_data.get("next")
semaphore = asyncio.Semaphore(self._playlist_concurrency)
_next = track_data.get("next")
if _next:
next_page_url = AM_BASE_URL + _next
while next_page_url is not None:
resp = await self.session.get(next_page_url, headers=self.headers)
async def fetch_page(url: str) -> List[Song]:
async with semaphore:
resp = await self.session.get(url, headers=self.headers)
if resp.status != 200:
if self._log:
self._log.warning(
f"Apple Music page fetch failed {resp.status} {resp.reason} for {url}",
)
return []
pj: dict = await resp.json(loads=json.loads)
songs = [Song(track) for track in pj.get("data", [])]
# Return songs; we will look for pj.get('next') in streaming iterator variant
return songs, pj.get("next") # type: ignore
raise AppleMusicRequestException(
f"Error while fetching results: {resp.status} {resp.reason}",
)
# We'll implement a wave-based approach similar to Spotify but need to follow cursors.
# Because we cannot know all cursors upfront, we'll iteratively fetch waves.
waves: List[List[Song]] = []
cursors: List[str] = []
if next_cursor:
cursors.append(next_cursor)
next_data: dict = await resp.json(loads=json.loads)
album_tracks.extend(Song(track) for track in next_data["data"])
# Limit total waves to avoid infinite loops in malformed responses
max_waves = 50
wave_size = self._playlist_concurrency * 2
wave_counter = 0
while cursors and wave_counter < max_waves:
current = cursors[:wave_size]
cursors = cursors[wave_size:]
tasks = [
fetch_page(AM_BASE_URL + cursor) for cursor in current # type: ignore[arg-type]
]
results = await asyncio.gather(*tasks, return_exceptions=True)
for res in results:
if isinstance(res, tuple): # (songs, next)
songs, nxt = res
if songs:
waves.append(songs)
if nxt:
cursors.append(nxt)
wave_counter += 1
for w in waves:
album_tracks.extend(w)
_next = next_data.get("next")
if _next:
next_page_url = AM_BASE_URL + _next
else:
next_page_url = None
return Playlist(data, album_tracks)
async def iter_playlist_tracks(
self,
*,
query: str,
batch_size: int = 100,
) -> AsyncGenerator[List[Song], None]:
"""Stream Apple Music playlist tracks in batches.
Parameters
----------
query: str
Apple Music playlist URL.
batch_size: int
Logical grouping size for yielded batches.
"""
if not self.token or datetime.utcnow() > self.expiry:
await self.request_token()
result = AM_URL_REGEX.match(query)
if not result or result.group("type") != "playlist":
raise InvalidAppleMusicURL("Provided query is not a valid Apple Music playlist URL.")
country = result.group("country")
playlist_id = result.group("id")
request_url = AM_REQ_URL.format(country=country, type="playlist", id=playlist_id)
resp = await self.session.get(request_url, headers=self.headers)
if resp.status != 200:
raise AppleMusicRequestException(
f"Error while fetching results: {resp.status} {resp.reason}",
)
data: dict = await resp.json(loads=json.loads)
playlist_data = data["data"][0]
track_data: dict = playlist_data["relationships"]["tracks"]
first_page_tracks = [Song(track) for track in track_data["data"]]
for i in range(0, len(first_page_tracks), batch_size):
yield first_page_tracks[i : i + batch_size]
next_cursor = track_data.get("next")
semaphore = asyncio.Semaphore(self._playlist_concurrency)
async def fetch(cursor: str) -> tuple[List[Song], Optional[str]]:
url = AM_BASE_URL + cursor
async with semaphore:
r = await self.session.get(url, headers=self.headers)
if r.status != 200:
if self._log:
self._log.warning(
f"Skipping Apple Music page due to {r.status} {r.reason}",
)
return [], None
pj: dict = await r.json(loads=json.loads)
songs = [Song(track) for track in pj.get("data", [])]
return songs, pj.get("next")
# Rolling waves of fetches following cursor chain
max_waves = 50
wave_size = self._playlist_concurrency * 2
waves = 0
cursors: List[str] = []
if next_cursor:
cursors.append(next_cursor)
while cursors and waves < max_waves:
current = cursors[:wave_size]
cursors = cursors[wave_size:]
results = await asyncio.gather(*[fetch(c) for c in current])
for songs, nxt in results:
if songs:
for j in range(0, len(songs), batch_size):
yield songs[j : j + batch_size]
if nxt:
cursors.append(nxt)
waves += 1

View File

@ -34,11 +34,6 @@ class SearchType(Enum):
ytsearch = "ytsearch"
ytmsearch = "ytmsearch"
scsearch = "scsearch"
other = "other"
@classmethod
def _missing_(cls, value: object) -> "SearchType": # type: ignore[override]
return cls.other
def __str__(self) -> str:
return self.value
@ -59,8 +54,6 @@ class TrackType(Enum):
TrackType.HTTP defines that the track is from an HTTP source.
TrackType.LOCAL defines that the track is from a local source.
TrackType.OTHER defines that the track is from an unknown source (possible from 3rd-party plugins).
"""
# We don't have to define anything special for these, since these just serve as flags
@ -70,11 +63,6 @@ class TrackType(Enum):
APPLE_MUSIC = "apple_music"
HTTP = "http"
LOCAL = "local"
OTHER = "other"
@classmethod
def _missing_(cls, value: object) -> "TrackType": # type: ignore[override]
return cls.OTHER
def __str__(self) -> str:
return self.value
@ -91,8 +79,6 @@ class PlaylistType(Enum):
PlaylistType.SPOTIFY defines that the playlist is from Spotify
PlaylistType.APPLE_MUSIC defines that the playlist is from Apple Music.
PlaylistType.OTHER defines that the playlist is from an unknown source (possible from 3rd-party plugins).
"""
# We don't have to define anything special for these, since these just serve as flags
@ -100,11 +86,6 @@ class PlaylistType(Enum):
SOUNDCLOUD = "soundcloud"
SPOTIFY = "spotify"
APPLE_MUSIC = "apple_music"
OTHER = "other"
@classmethod
def _missing_(cls, value: object) -> "PlaylistType": # type: ignore[override]
return cls.OTHER
def __str__(self) -> str:
return self.value
@ -218,12 +199,8 @@ class URLRegex:
"""
# Spotify share links can include query parameters like ?si=XXXX, a trailing slash,
# or an intl locale segment (e.g. /intl-en/). Broaden the regex so we still capture
# the type and id while ignoring extra parameters. This prevents the URL from being
# treated as a generic Lavalink identifier and ensures internal Spotify handling runs.
SPOTIFY_URL = re.compile(
r"https?://open\.spotify\.com/(?:intl-[a-zA-Z-]+/)?(?P<type>album|playlist|track|artist)/(?P<id>[a-zA-Z0-9]+)(?:/)?(?:\?.*)?$",
r"https?://open.spotify.com/(?P<type>album|playlist|track|artist)/(?P<id>[a-zA-Z0-9]+)",
)
DISCORD_MP3_URL = re.compile(
@ -244,17 +221,14 @@ class URLRegex:
r"(?P<video>^.*?)(\?t|&start)=(?P<time>\d+)?.*",
)
# Apple Music links sometimes append additional query parameters (e.g. &l=en, &uo=4).
# Allow arbitrary query parameters so valid links are captured and parsed.
AM_URL = re.compile(
r"https?://music\.apple\.com/(?P<country>[a-zA-Z]{2})/"
r"(?P<type>album|playlist|song|artist)/(?P<name>.+?)/(?P<id>[^/?]+?)(?:/)?(?:\?.*)?$",
r"https?://music.apple.com/(?P<country>[a-zA-Z]{2})/"
r"(?P<type>album|playlist|song|artist)/(?P<name>.+)/(?P<id>[^?]+)",
)
# Single-in-album links may also carry extra query params beyond the ?i=<trackid> token.
AM_SINGLE_IN_ALBUM_REGEX = re.compile(
r"https?://music\.apple\.com/(?P<country>[a-zA-Z]{2})/(?P<type>album|playlist|song|artist)/"
r"(?P<name>.+)/(?P<id>[^/?]+)(\?i=)(?P<id2>[^&]+)(?:&.*)?$",
r"https?://music.apple.com/(?P<country>[a-zA-Z]{2})/(?P<type>album|playlist|song|artist)/"
r"(?P<name>.+)/(?P<id>.+)(\?i=)(?P<id2>.+)",
)
SOUNDCLOUD_URL = re.compile(
@ -299,10 +273,3 @@ class LogLevel(IntEnum):
WARN = 30
ERROR = 40
CRITICAL = 50
@classmethod
def from_str(cls, level_str):
try:
return cls[level_str.upper()]
except KeyError:
raise ValueError(f"No such log level: {level_str}")

View File

@ -128,6 +128,7 @@ class TrackExceptionEvent(PomiceEvent):
def __init__(self, data: dict, player: Player):
self.player: Player = player
assert self.player._ending_track is not None
self.track: Optional[Track] = self.player._ending_track
# Error is for Lavalink <= 3.3
self.exception: str = data.get(

View File

@ -77,12 +77,6 @@ class Equalizer(Filter):
def __repr__(self) -> str:
return f"<Pomice.EqualizerFilter tag={self.tag} eq={self.eq} raw={self.raw}>"
def __eq__(self, __value: object) -> bool:
if not isinstance(__value, Equalizer):
return False
return self.raw == __value.raw
@classmethod
def flat(cls) -> "Equalizer":
"""Equalizer preset which represents a flat EQ board,
@ -237,16 +231,6 @@ class Timescale(Filter):
def __repr__(self) -> str:
return f"<Pomice.TimescaleFilter tag={self.tag} speed={self.speed} pitch={self.pitch} rate={self.rate}>"
def __eq__(self, __value: object) -> bool:
if not isinstance(__value, Timescale):
return False
return (
self.speed == __value.speed
and self.pitch == __value.pitch
and self.rate == __value.rate
)
class Karaoke(Filter):
"""Filter which filters the vocal track from any song and leaves the instrumental.
@ -286,17 +270,6 @@ class Karaoke(Filter):
f"filter_band={self.filter_band} filter_width={self.filter_width}>"
)
def __eq__(self, __value: object) -> bool:
if not isinstance(__value, Karaoke):
return False
return (
self.level == __value.level
and self.mono_level == __value.mono_level
and self.filter_band == __value.filter_band
and self.filter_width == __value.filter_width
)
class Tremolo(Filter):
"""Filter which produces a wavering tone in the music,
@ -332,12 +305,6 @@ class Tremolo(Filter):
f"<Pomice.TremoloFilter tag={self.tag} frequency={self.frequency} depth={self.depth}>"
)
def __eq__(self, __value: object) -> bool:
if not isinstance(__value, Tremolo):
return False
return self.frequency == __value.frequency and self.depth == __value.depth
class Vibrato(Filter):
"""Filter which produces a wavering tone in the music, similar to the Tremolo filter,
@ -373,12 +340,6 @@ class Vibrato(Filter):
f"<Pomice.VibratoFilter tag={self.tag} frequency={self.frequency} depth={self.depth}>"
)
def __eq__(self, __value: object) -> bool:
if not isinstance(__value, Vibrato):
return False
return self.frequency == __value.frequency and self.depth == __value.depth
class Rotation(Filter):
"""Filter which produces a stereo-like panning effect, which sounds like
@ -396,12 +357,6 @@ class Rotation(Filter):
def __repr__(self) -> str:
return f"<Pomice.RotationFilter tag={self.tag} rotation_hertz={self.rotation_hertz}>"
def __eq__(self, __value: object) -> bool:
if not isinstance(__value, Rotation):
return False
return self.rotation_hertz == __value.rotation_hertz
class ChannelMix(Filter):
"""Filter which manually adjusts the panning of the audio, which can make
@ -463,17 +418,6 @@ class ChannelMix(Filter):
f"right_to_left={self.right_to_left} right_to_right={self.right_to_right}>"
)
def __eq__(self, __value: object) -> bool:
if not isinstance(__value, ChannelMix):
return False
return (
self.left_to_left == __value.left_to_left
and self.left_to_right == __value.left_to_right
and self.right_to_left == __value.right_to_left
and self.right_to_right == __value.right_to_right
)
class Distortion(Filter):
"""Filter which generates a distortion effect. Useful for certain filter implementations where
@ -535,21 +479,6 @@ class Distortion(Filter):
f"tan_scale={self.tan_scale} offset={self.offset} scale={self.scale}"
)
def __eq__(self, __value: object) -> bool:
if not isinstance(__value, Distortion):
return False
return (
self.sin_offset == __value.sin_offset
and self.sin_scale == __value.sin_scale
and self.cos_offset == __value.cos_offset
and self.cos_scale == __value.cos_scale
and self.tan_offset == __value.tan_offset
and self.tan_scale == __value.tan_scale
and self.offset == __value.offset
and self.scale == __value.scale
)
class LowPass(Filter):
"""Filter which supresses higher frequencies and allows lower frequencies to pass.
@ -566,9 +495,3 @@ class LowPass(Filter):
def __repr__(self) -> str:
return f"<Pomice.LowPass tag={self.tag} smoothing={self.smoothing}>"
def __eq__(self, __value: object) -> bool:
if not isinstance(__value, LowPass):
return False
return self.smoothing == __value.smoothing

View File

@ -98,6 +98,9 @@ class Track:
if not isinstance(other, Track):
return False
if self.ctx and other.ctx:
return other.track_id == self.track_id and other.ctx.message.id == self.ctx.message.id
return other.track_id == self.track_id
def __str__(self) -> str:

View File

@ -79,27 +79,6 @@ class Filters:
if filter.tag == filter_tag:
del self._filters[index]
def edit_filter(self, *, filter_tag: str, to_apply: Filter) -> None:
"""Edits a filter in the list of filters applied using its filter tag and replaces it with the new filter."""
if not any(f for f in self._filters if f.tag == filter_tag):
raise FilterTagInvalid("A filter with that tag was not found.")
for index, filter in enumerate(self._filters):
if filter.tag == filter_tag:
if not type(filter) == type(to_apply):
raise FilterInvalidArgument(
"Edited filter is not the same type as the current filter.",
)
if self._filters[index] == to_apply:
raise FilterInvalidArgument("Edited filter is the same as the current filter.")
if to_apply.tag != filter_tag:
raise FilterInvalidArgument(
"Edited filter tag is not the same as the current filter tag.",
)
self._filters[index] = to_apply
def has_filter(self, *, filter_tag: str) -> bool:
"""Checks if a filter exists in the list of filters using its filter tag"""
return any(f for f in self._filters if f.tag == filter_tag)
@ -213,7 +192,7 @@ class Player(VoiceProtocol):
difference = (time.time() * 1000) - self._last_update
position = self._last_position + difference
return round(min(position, current.length))
return min(position, current.length)
@property
def rate(self) -> float:
@ -298,8 +277,7 @@ class Player(VoiceProtocol):
self._last_update = int(state.get("time", 0))
self._is_connected = bool(state.get("connected"))
self._last_position = int(state.get("position", 0))
if self._log:
self._log.debug(f"Got player update state with data {state}")
self._log.debug(f"Got player update state with data {state}")
async def _dispatch_voice_update(self, voice_data: Optional[Dict[str, Any]] = None) -> None:
if {"sessionId", "event"} != self._voice_state.keys():
@ -320,10 +298,7 @@ class Player(VoiceProtocol):
data={"voice": data},
)
if self._log:
self._log.debug(
f"Dispatched voice update to {state['event']['endpoint']} with data {data}",
)
self._log.debug(f"Dispatched voice update to {state['event']['endpoint']} with data {data}")
async def on_voice_server_update(self, data: VoiceServerUpdate) -> None:
self._voice_state.update({"event": data})
@ -339,10 +314,6 @@ class Player(VoiceProtocol):
return
channel = self.guild.get_channel(int(channel_id))
if self.channel != channel:
self.channel = channel
if not channel:
await self.disconnect()
self._voice_state.clear()
@ -357,7 +328,7 @@ class Player(VoiceProtocol):
event_type: str = data["type"]
event: PomiceEvent = getattr(events, event_type)(data, self)
if isinstance(event, TrackEndEvent) and event.reason not in ("REPLACED", "replaced"):
if isinstance(event, TrackEndEvent) and event.reason != "REPLACED":
self._current = None
event.dispatch(self._bot)
@ -365,8 +336,7 @@ class Player(VoiceProtocol):
if isinstance(event, TrackStartEvent):
self._ending_track = self._current
if self._log:
self._log.debug(f"Dispatched event {data['type']} to player.")
self._log.debug(f"Dispatched event {data['type']} to player.")
async def _refresh_endpoint_uri(self, session_id: Optional[str]) -> None:
self._player_endpoint_uri = f"sessions/{session_id}/players"
@ -388,15 +358,14 @@ class Player(VoiceProtocol):
data=data or None,
)
if self._log:
self._log.debug(f"Swapped all players to new node {new_node._identifier}.")
self._log.debug(f"Swapped all players to new node {new_node._identifier}.")
async def get_tracks(
self,
query: str,
*,
ctx: Optional[commands.Context] = None,
search_type: SearchType | None = SearchType.ytsearch,
search_type: SearchType = SearchType.ytsearch,
filters: Optional[List[Filter]] = None,
) -> Optional[Union[List[Track], Playlist]]:
"""Fetches tracks from the node's REST api to parse into Lavalink.
@ -413,21 +382,8 @@ class Player(VoiceProtocol):
"""
return await self._node.get_tracks(query, ctx=ctx, search_type=search_type, filters=filters)
async def build_track(self, identifier: str, ctx: Optional[commands.Context] = None) -> Track:
"""
Builds a track using a valid track identifier
You can also pass in a discord.py Context object to get a
Context object on the track it builds.
"""
return await self._node.build_track(identifier, ctx=ctx)
async def get_recommendations(
self,
*,
track: Track,
ctx: Optional[commands.Context] = None,
self, *, track: Track, ctx: Optional[commands.Context] = None
) -> Optional[Union[List[Track], Playlist]]:
"""
Gets recommendations from either YouTube or Spotify.
@ -437,12 +393,7 @@ class Player(VoiceProtocol):
return await self._node.get_recommendations(track=track, ctx=ctx)
async def connect(
self,
*,
timeout: float,
reconnect: bool,
self_deaf: bool = False,
self_mute: bool = False,
self, *, timeout: float, reconnect: bool, self_deaf: bool = False, self_mute: bool = False
) -> None:
await self.guild.change_voice_state(
channel=self.channel,
@ -462,8 +413,7 @@ class Player(VoiceProtocol):
data={"encodedTrack": None},
)
if self._log:
self._log.debug(f"Player has been stopped.")
self._log.debug(f"Player has been stopped.")
async def disconnect(self, *, force: bool = False) -> None:
"""Disconnects the player from voice."""
@ -481,34 +431,24 @@ class Player(VoiceProtocol):
except AttributeError:
# 'NoneType' has no attribute '_get_voice_client_key' raised by self.cleanup() ->
# assume we're already disconnected and cleaned up
assert self.channel is None and not self.is_connected
assert not self.is_connected and not self.channel
self._node._players.pop(self.guild.id)
if self.node.is_connected:
await self._node.send(
method="DELETE",
path=self._player_endpoint_uri,
guild_id=self._guild.id,
)
await self._node.send(
method="DELETE",
path=self._player_endpoint_uri,
guild_id=self._guild.id,
)
if self._log:
self._log.debug("Player has been destroyed.")
self._log.debug("Player has been destroyed.")
async def play(
self,
track: Track,
*,
start: int = 0,
end: int = 0,
ignore_if_playing: bool = False,
self, track: Track, *, start: int = 0, end: int = 0, ignore_if_playing: bool = False
) -> Track:
"""Plays a track. If a Spotify track is passed in, it will be handled accordingly."""
if not track._search_type:
track.original = track
# Make sure we've never searched the track before
if track._search_type and track.original is None:
if track.original is None:
# First lets try using the tracks ISRC, every track has one (hopefully)
try:
if not track.isrc:
@ -588,10 +528,9 @@ class Player(VoiceProtocol):
query=f"noReplace={ignore_if_playing}",
)
if self._log:
self._log.debug(
f"Playing {track.title} from uri {track.uri} with a length of {track.length}",
)
self._log.debug(
f"Playing {track.title} from uri {track.uri} with a length of {track.length}",
)
return self._current
@ -612,8 +551,7 @@ class Player(VoiceProtocol):
data={"position": position},
)
if self._log:
self._log.debug(f"Seeking to {position}.")
self._log.debug(f"Seeking to {position}.")
return self.position
async def set_pause(self, pause: bool) -> bool:
@ -626,8 +564,7 @@ class Player(VoiceProtocol):
)
self._paused = pause
if self._log:
self._log.debug(f"Player has been {'paused' if pause else 'resumed'}.")
self._log.debug(f"Player has been {'paused' if pause else 'resumed'}.")
return self._paused
async def set_volume(self, volume: int) -> int:
@ -640,8 +577,7 @@ class Player(VoiceProtocol):
)
self._volume = volume
if self._log:
self._log.debug(f"Player volume has been adjusted to {volume}")
self._log.debug(f"Player volume has been adjusted to {volume}")
return self._volume
async def move_to(self, channel: VoiceChannel) -> None:
@ -670,11 +606,9 @@ class Player(VoiceProtocol):
data={"filters": payload},
)
if self._log:
self._log.debug(f"Filter has been applied to player with tag {_filter.tag}")
self._log.debug(f"Filter has been applied to player with tag {_filter.tag}")
if fast_apply:
if self._log:
self._log.debug(f"Fast apply passed, now applying filter instantly.")
self._log.debug(f"Fast apply passed, now applying filter instantly.")
await self.seek(self.position)
return self._filters
@ -695,44 +629,9 @@ class Player(VoiceProtocol):
guild_id=self._guild.id,
data={"filters": payload},
)
if self._log:
self._log.debug(f"Filter has been removed from player with tag {filter_tag}")
self._log.debug(f"Filter has been removed from player with tag {filter_tag}")
if fast_apply:
if self._log:
self._log.debug(f"Fast apply passed, now removing filter instantly.")
await self.seek(self.position)
return self._filters
async def edit_filter(
self,
*,
filter_tag: str,
edited_filter: Filter,
fast_apply: bool = False,
) -> Filters:
"""Edits a filter from the player using its filter tag and a new filter of the same type.
The filter to be replaced must have the same tag as the one you are replacing it with.
This will only work if you are using a version of Lavalink that supports filters.
If you would like for the filter to apply instantly, set the `fast_apply` arg to `True`.
(You must have a song playing in order for `fast_apply` to work.)
"""
self._filters.edit_filter(filter_tag=filter_tag, to_apply=edited_filter)
payload = self._filters.get_all_payloads()
await self._node.send(
method="PATCH",
path=self._player_endpoint_uri,
guild_id=self._guild.id,
data={"filters": payload},
)
if self._log:
self._log.debug(f"Filter with tag {filter_tag} has been edited to {edited_filter!r}")
if fast_apply:
if self._log:
self._log.debug(f"Fast apply passed, now editing filter instantly.")
self._log.debug(f"Fast apply passed, now removing filter instantly.")
await self.seek(self.position)
return self._filters
@ -756,10 +655,8 @@ class Player(VoiceProtocol):
guild_id=self._guild.id,
data={"filters": {}},
)
if self._log:
self._log.debug(f"All filters have been removed from player.")
self._log.debug(f"All filters have been removed from player.")
if fast_apply:
if self._log:
self._log.debug(f"Fast apply passed, now removing all filters instantly.")
self._log.debug(f"Fast apply passed, now removing all filters instantly.")
await self.seek(self.position)

View File

@ -17,24 +17,16 @@ from typing import Union
from urllib.parse import quote
import aiohttp
import orjson as json
from discord import Client
from discord.ext import commands
from discord.utils import MISSING
try:
from websockets.legacy import client # websockets >= 10.0
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 LogLevel
from .exceptions import AppleMusicNotEnabled
from .exceptions import InvalidSpotifyClientAuthorization
from .exceptions import LavalinkVersionIncompatible
from .exceptions import NodeConnectionFailure
@ -79,8 +71,6 @@ class Node:
"_password",
"_identifier",
"_heartbeat",
"_resume_key",
"_resume_timeout",
"_secure",
"_fallback",
"_log_level",
@ -101,6 +91,7 @@ class Node:
"_apple_music_client",
"_route_planner",
"_log",
"_log_handler",
"_stats",
"available",
)
@ -115,16 +106,15 @@ class Node:
password: str,
identifier: str,
secure: bool = False,
heartbeat: int = 120,
resume_key: Optional[str] = None,
resume_timeout: int = 60,
heartbeat: int = 30,
loop: Optional[asyncio.AbstractEventLoop] = None,
session: Optional[aiohttp.ClientSession] = None,
spotify_client_id: Optional[str] = None,
spotify_client_secret: Optional[str] = None,
apple_music: bool = False,
fallback: bool = False,
logger: Optional[logging.Logger] = None,
log_level: LogLevel = LogLevel.INFO,
log_handler: Optional[logging.Handler] = None,
):
if not isinstance(port, int):
raise TypeError("Port must be an integer")
@ -136,17 +126,17 @@ class Node:
self._password: str = password
self._identifier: str = identifier
self._heartbeat: int = heartbeat
self._resume_key: Optional[str] = resume_key
self._resume_timeout: int = resume_timeout
self._secure: bool = secure
self._fallback: bool = fallback
self._log_level: LogLevel = log_level
self._log_handler = log_handler
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}"
self._session: aiohttp.ClientSession = session # type: ignore
self._loop: asyncio.AbstractEventLoop = loop or asyncio.get_event_loop()
self._websocket: client.WebSocketClientProtocol
self._websocket: aiohttp.ClientWebSocketResponse
self._task: asyncio.Task = None # type: ignore
self._session_id: Optional[str] = None
@ -154,7 +144,7 @@ class Node:
self._version: LavalinkVersion = LavalinkVersion(0, 0, 0)
self._route_planner = RoutePlanner(self)
self._log = logger
self._log = self._setup_logging(self._log_level)
if not self._bot.user:
raise NodeCreationError("Bot user is not ready yet.")
@ -215,7 +205,7 @@ class Node:
@property
def player_count(self) -> int:
"""Property which returns how many players are connected to this node"""
return len(self.players.values())
return len(self.players)
@property
def pool(self) -> Type[NodePool]:
@ -232,6 +222,29 @@ class Node:
"""Alias for `Node.latency`, returns the latency of the node"""
return self.latency
def _setup_logging(self, level: LogLevel) -> logging.Logger:
logger = logging.getLogger("pomice")
logger.setLevel(level)
handler = None
if self._log_handler:
handler = self._log_handler
else:
handler = logging.StreamHandler()
dt_fmt = "%Y-%m-%d %H:%M:%S"
formatter = logging.Formatter(
"[{asctime}] [{levelname:<8}] {name}: {message}",
dt_fmt,
style="{",
)
handler.setFormatter(formatter)
if handler:
logger.handlers.clear()
logger.addHandler(handler)
return logger
async def _handle_version_check(self, version: str) -> None:
if version.endswith("-SNAPSHOT"):
# we're just gonna assume all snapshot versions correlate with v4
@ -253,8 +266,7 @@ class Node:
int(_version_groups[2] or 0),
)
if self._log:
self._log.debug(f"Parsed Lavalink version: {major}.{minor}.{fix}")
self._log.debug(f"Parsed Lavalink version: {major}.{minor}.{fix}")
self._version = LavalinkVersion(major=major, minor=minor, fix=fix)
if self._version < LavalinkVersion(3, 7, 0):
self._available = False
@ -304,59 +316,25 @@ class Node:
await self.disconnect()
async def _configure_resuming(self) -> None:
if not self._resume_key:
return
data: Dict[str, Union[int, str, bool]] = {"timeout": self._resume_timeout}
if self._version.major == 3:
data["resumingKey"] = self._resume_key
elif self._version.major == 4:
if self._log:
self._log.warning("Using a resume key with Lavalink v4 is deprecated.")
data["resuming"] = True
await self.send(
method="PATCH",
path=f"sessions/{self._session_id}",
include_version=True,
data=data,
)
async def _listen(self) -> None:
backoff = ExponentialBackoff(base=7)
while True:
try:
msg = await self._websocket.recv()
data = json.loads(msg)
if self._log:
self._log.debug(f"Recieved raw websocket message {msg}")
self._loop.create_task(self._handle_ws_msg(data=data))
except exceptions.ConnectionClosed:
if self.player_count > 0:
for _player in self.players.values():
self._loop.create_task(_player.destroy())
msg = await self._websocket.receive()
if msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.CLOSING):
if self._fallback:
self._loop.create_task(self._handle_node_switch())
self._loop.create_task(self._websocket.close())
backoff = ExponentialBackoff(base=7)
await self._handle_node_switch()
retry = backoff.delay()
if self._log:
self._log.debug(
f"Retrying connection to Node {self._identifier} in {retry} secs",
)
await asyncio.sleep(retry)
if not self.is_connected:
self._loop.create_task(self.connect(reconnect=True))
else:
self._loop.create_task(self._handle_payload(msg.json()))
async def _handle_ws_msg(self, data: dict) -> None:
if self._log:
self._log.debug(f"Recieved raw payload from Node {self._identifier} with data {data}")
async def _handle_payload(self, data: dict) -> None:
op = data.get("op", None)
if not op:
return
if op == "stats":
self._stats = NodeStats(data)
@ -364,20 +342,21 @@ class Node:
if op == "ready":
self._session_id = data["sessionId"]
await self._configure_resuming()
if not "guildId" in data:
return
player: Optional[Player] = self._players.get(int(data["guildId"]))
player = self._players.get(int(data["guildId"]))
if not player:
return
if op == "event":
return await player._dispatch_event(data)
await player._dispatch_event(data)
return
if op == "playerUpdate":
return await player._update_state(data)
await player._update_state(data)
return
async def send(
self,
@ -408,10 +387,9 @@ class Node:
headers=self._headers,
json=data or {},
)
if self._log:
self._log.debug(
f"Making REST request to Node {self._identifier} with method {method} to {uri}",
)
self._log.debug(
f"Making REST request to Node {self._identifier} with method {method} to {uri}",
)
if resp.status >= 300:
resp_data: dict = await resp.json()
raise NodeRestException(
@ -419,47 +397,34 @@ class Node:
)
if method == "DELETE" or resp.status == 204:
if self._log:
self._log.debug(
f"REST request to Node {self._identifier} with method {method} to {uri} completed sucessfully and returned no data.",
)
self._log.debug(
f"REST request to Node {self._identifier} with method {method} to {uri} completed sucessfully and returned no data.",
)
return await resp.json(content_type=None)
if resp.content_type == "text/plain":
if self._log:
self._log.debug(
f"REST request to Node {self._identifier} with method {method} to {uri} completed sucessfully and returned text with body {await resp.text()}",
)
self._log.debug(
f"REST request to Node {self._identifier} with method {method} to {uri} completed sucessfully and returned text with body {await resp.text()}",
)
return await resp.text()
if self._log:
self._log.debug(
f"REST request to Node {self._identifier} with method {method} to {uri} completed sucessfully and returned JSON with body {await resp.json()}",
)
self._log.debug(
f"REST request to Node {self._identifier} with method {method} to {uri} completed sucessfully and returned JSON with body {await resp.json()}",
)
return await resp.json()
def get_player(self, guild_id: int) -> Optional[Player]:
"""Takes a guild ID as a parameter. Returns a pomice Player object or None."""
return self._players.get(guild_id, None)
async def connect(self, *, reconnect: bool = False) -> Node:
async def connect(self, *, reconnect: bool = False) -> "Node":
"""Initiates a connection with a Lavalink node and adds it to the node pool."""
await self._bot.wait_until_ready()
start = time.perf_counter()
if not self._session:
# Configure connection pooling for optimal concurrent request performance
connector = aiohttp.TCPConnector(
limit=100, # Total connection limit
limit_per_host=30, # Per-host connection limit
ttl_dns_cache=300, # DNS cache TTL in seconds
)
timeout = aiohttp.ClientTimeout(total=30, connect=10)
self._session = aiohttp.ClientSession(
connector=connector,
timeout=timeout,
)
self._session = aiohttp.ClientSession()
try:
if not reconnect:
@ -473,28 +438,23 @@ class Node:
await self._handle_version_check(version=version)
await self._set_ext_client_session(session=self._session)
if self._log:
self._log.debug(
f"Version check from Node {self._identifier} successful. Returned version {version}",
)
self._log.debug(
f"Version check from Node {self._identifier} successful. Returned version {version}",
)
self._websocket = await client.connect( # type: ignore
self._websocket = await self._session.ws_connect(
f"{self._websocket_uri}/v{self._version.major}/websocket",
extra_headers=self._headers,
ping_interval=self._heartbeat,
headers=self._headers,
heartbeat=self._heartbeat,
)
if reconnect:
if self._log:
self._log.debug(f"Trying to reconnect to Node {self._identifier}...")
if self.player_count:
for player in self.players.values():
await player._refresh_endpoint_uri(self._session_id)
for player in self.players.values():
await player._refresh_endpoint_uri(self._session_id)
if self._log:
self._log.debug(
f"Node {self._identifier} successfully connected to websocket using {self._websocket_uri}/v{self._version.major}/websocket",
)
self._log.debug(
f"Node {self._identifier} successfully connected to websocket using {self._websocket_uri}/v{self._version.major}/websocket",
)
if not self._task:
self._task = self._loop.create_task(self._listen())
@ -503,19 +463,18 @@ class Node:
end = time.perf_counter()
if self._log:
self._log.info(f"Connected to node {self._identifier}. Took {end - start:.3f}s")
self._log.info(f"Connected to node {self._identifier}. Took {end - start:.3f}s")
return self
except (aiohttp.ClientConnectorError, OSError, ConnectionRefusedError):
except (aiohttp.ClientConnectorError, ConnectionRefusedError):
raise NodeConnectionFailure(
f"The connection to node '{self._identifier}' failed.",
) from None
except exceptions.InvalidHandshake:
except aiohttp.WSServerHandshakeError:
raise NodeConnectionFailure(
f"The password for node '{self._identifier}' is invalid.",
) from None
except exceptions.InvalidURI:
except aiohttp.InvalidURL:
raise NodeConnectionFailure(
f"The URI for node '{self._identifier}' is invalid.",
) from None
@ -529,23 +488,20 @@ class Node:
for player in self.players.copy().values():
await player.destroy()
if self._log:
self._log.debug("All players disconnected from node.")
self._log.debug("All players disconnected from node.")
await self._websocket.close()
await self._session.close()
if self._log:
self._log.debug("Websocket and http session closed.")
self._log.debug("Websocket and http session closed.")
del self._pool._nodes[self._identifier]
self.available = False
self._task.cancel()
end = time.perf_counter()
if self._log:
self._log.info(
f"Successfully disconnected from node {self._identifier} and closed all sessions. Took {end - start:.3f}s",
)
self._log.info(
f"Successfully disconnected from node {self._identifier} and closed all sessions. Took {end - start:.3f}s",
)
async def build_track(self, identifier: str, ctx: Optional[commands.Context] = None) -> Track:
"""
@ -558,16 +514,13 @@ class Node:
data: dict = await self.send(
method="GET",
path="decodetrack",
query=f"encodedTrack={quote(identifier)}",
query=f"encodedTrack={identifier}",
)
track_info = data["info"] if self._version.major >= 4 else data
return Track(
track_id=identifier,
ctx=ctx,
info=track_info,
track_type=TrackType(track_info["sourceName"]),
info=data,
track_type=TrackType(data["sourceName"]),
)
async def get_tracks(
@ -575,7 +528,7 @@ class Node:
query: str,
*,
ctx: Optional[commands.Context] = None,
search_type: Optional[SearchType] = SearchType.ytsearch,
search_type: SearchType = SearchType.ytsearch,
filters: Optional[List[Filter]] = None,
) -> Optional[Union[Playlist, List[Track]]]:
"""Fetches tracks from the node's REST api to parse into Lavalink.
@ -596,13 +549,13 @@ class Node:
for filter in filters:
filter.set_preload()
# Due to the inclusion of plugins in the v4 update
# we are doing away with raising an error if pomice detects
# either a Spotify or Apple Music URL and the respective client
# is not enabled. Instead, we will just only parse the URL
# if the client is enabled and the URL is valid.
if URLRegex.AM_URL.match(query):
if not self._apple_music_client:
raise AppleMusicNotEnabled(
"You must have Apple Music functionality enabled in order to play Apple Music tracks."
"Please set apple_music to True in your Node class.",
)
if self._apple_music_client and URLRegex.AM_URL.match(query):
apple_music_results = await self._apple_music_client.search(query=query)
if isinstance(apple_music_results, applemusic.Song):
return [
@ -610,7 +563,7 @@ class Node:
track_id=apple_music_results.id,
ctx=ctx,
track_type=TrackType.APPLE_MUSIC,
search_type=search_type or SearchType.ytsearch,
search_type=search_type,
filters=filters,
info={
"title": apple_music_results.name,
@ -632,7 +585,7 @@ class Node:
track_id=track.id,
ctx=ctx,
track_type=TrackType.APPLE_MUSIC,
search_type=search_type or SearchType.ytsearch,
search_type=search_type,
filters=filters,
info={
"title": track.name,
@ -661,7 +614,14 @@ class Node:
uri=apple_music_results.url,
)
elif self._spotify_client and URLRegex.SPOTIFY_URL.match(query):
elif URLRegex.SPOTIFY_URL.match(query):
if not self._spotify_client_id and not self._spotify_client_secret:
raise InvalidSpotifyClientAuthorization(
"You did not provide proper Spotify client authorization credentials. "
"If you would like to use the Spotify searching feature, "
"please obtain Spotify API credentials here: https://developer.spotify.com/",
)
spotify_results = await self._spotify_client.search(query=query) # type: ignore
if isinstance(spotify_results, spotify.Track):
@ -670,7 +630,7 @@ class Node:
track_id=spotify_results.id,
ctx=ctx,
track_type=TrackType.SPOTIFY,
search_type=search_type or SearchType.ytsearch,
search_type=search_type,
filters=filters,
info={
"title": spotify_results.name,
@ -692,7 +652,7 @@ class Node:
track_id=track.id,
ctx=ctx,
track_type=TrackType.SPOTIFY,
search_type=search_type or SearchType.ytsearch,
search_type=search_type,
filters=filters,
info={
"title": track.name,
@ -721,14 +681,63 @@ class Node:
uri=spotify_results.uri,
)
elif discord_url := URLRegex.DISCORD_MP3_URL.match(query):
data: dict = await self.send(
method="GET",
path="loadtracks",
query=f"identifier={quote(query)}",
)
track: dict = data["tracks"][0]
info: dict = track["info"]
return [
Track(
track_id=track["track"],
info={
"title": discord_url.group("file"),
"author": "Unknown",
"length": info["length"],
"uri": info["uri"],
"position": info["position"],
"identifier": info["identifier"],
},
ctx=ctx,
track_type=TrackType.HTTP,
filters=filters,
),
]
elif path.exists(path.dirname(query)):
local_file = Path(query)
data: dict = await self.send( # type: ignore
method="GET",
path="loadtracks",
query=f"identifier={quote(query)}",
)
track: dict = data["tracks"][0] # type: ignore
info: dict = track["info"] # type: ignore
return [
Track(
track_id=track["track"],
info={
"title": local_file.name,
"author": "Unknown",
"length": info["length"],
"uri": quote(local_file.as_uri()),
"position": info["position"],
"identifier": info["identifier"],
},
ctx=ctx,
track_type=TrackType.LOCAL,
filters=filters,
),
]
else:
if (
search_type
and not URLRegex.BASE_URL.match(query)
and not re.match(r"(?:[a-z]+?)search:.", query)
and not URLRegex.DISCORD_MP3_URL.match(query)
and not path.exists(path.dirname(query))
):
if not URLRegex.BASE_URL.match(query) and not re.match(r"(?:ytm?|sc)search:.", query):
query = f"{search_type}:{query}"
# If YouTube url contains a timestamp, capture it for use later.
@ -744,31 +753,21 @@ class Node:
load_type = data.get("loadType")
# Lavalink v4 changed the name of the key from "tracks" to "data"
# so lets account for that
data_type = "data" if self._version.major >= 4 else "tracks"
if not load_type:
raise TrackLoadError(
"There was an error while trying to load this track.",
)
elif load_type in ("LOAD_FAILED", "error"):
exception = data["data"] if self._version.major >= 4 else data["exception"]
elif load_type == "LOAD_FAILED":
exception = data["exception"]
raise TrackLoadError(
f"{exception['message']} [{exception['severity']}]",
)
elif load_type in ("NO_MATCHES", "empty"):
elif load_type == "NO_MATCHES":
return None
elif load_type in ("PLAYLIST_LOADED", "playlist"):
if self._version.major >= 4:
track_list = data[data_type]["tracks"]
playlist_info = data[data_type]["info"]
else:
track_list = data[data_type]
playlist_info = data["playlistInfo"]
elif load_type == "PLAYLIST_LOADED":
tracks = [
Track(
track_id=track["encoded"],
@ -776,60 +775,17 @@ class Node:
ctx=ctx,
track_type=TrackType(track["info"]["sourceName"]),
)
for track in track_list
for track in data["tracks"]
]
return Playlist(
playlist_info=playlist_info,
playlist_info=data["playlistInfo"],
tracks=tracks,
playlist_type=PlaylistType(tracks[0].track_type.value),
thumbnail=tracks[0].thumbnail,
uri=query,
)
elif load_type in ("SEARCH_RESULT", "TRACK_LOADED", "track", "search"):
if self._version.major >= 4 and isinstance(data[data_type], dict):
data[data_type] = [data[data_type]]
if path.exists(path.dirname(query)):
local_file = Path(query)
return [
Track(
track_id=track["encoded"],
info={
"title": local_file.name,
"author": "Unknown",
"length": track["info"]["length"],
"uri": quote(local_file.as_uri()),
"position": track["info"]["position"],
"identifier": track["info"]["identifier"],
},
ctx=ctx,
track_type=TrackType.LOCAL,
filters=filters,
)
for track in data[data_type]
]
elif discord_url := URLRegex.DISCORD_MP3_URL.match(query):
return [
Track(
track_id=track["encoded"],
info={
"title": discord_url.group("file"),
"author": "Unknown",
"length": track["info"]["length"],
"uri": track["info"]["uri"],
"position": track["info"]["position"],
"identifier": track["info"]["identifier"],
},
ctx=ctx,
track_type=TrackType.HTTP,
filters=filters,
)
for track in data[data_type]
]
elif load_type == "SEARCH_RESULT" or load_type == "TRACK_LOADED":
return [
Track(
track_id=track["encoded"],
@ -839,7 +795,7 @@ class Node:
filters=filters,
timestamp=timestamp,
)
for track in data[data_type]
for track in data["tracks"]
]
else:
@ -848,10 +804,7 @@ class Node:
)
async def get_recommendations(
self,
*,
track: Track,
ctx: Optional[commands.Context] = None,
self, *, track: Track, ctx: Optional[commands.Context] = None
) -> Optional[Union[List[Track], Playlist]]:
"""
Gets recommendations from either YouTube or Spotify.
@ -896,57 +849,6 @@ class Node:
"The specfied track must be either a YouTube or Spotify track to recieve recommendations.",
)
async def search_spotify_recommendations(
self,
query: str,
*,
ctx: Optional[commands.Context] = None,
filters: Optional[List[Filter]] = None,
) -> Optional[Union[List[Track], Playlist]]:
"""
Searches for recommendations on Spotify and returns a list of tracks based on the query.
You must have Spotify enabled for this to work.
You can pass in a discord.py Context object to get a
Context object on all tracks that get recommended.
"""
if not self._spotify_client:
raise InvalidSpotifyClientAuthorization(
"You must have Spotify enabled to use this feature.",
)
results = await self._spotify_client.track_search(query=query) # type: ignore
if not results:
raise TrackLoadError(
"Unable to find any tracks based on the query.",
)
tracks = [
Track(
track_id=track.id,
ctx=ctx,
track_type=TrackType.SPOTIFY,
info={
"title": track.name,
"author": track.artists,
"length": track.length,
"identifier": track.id,
"uri": track.uri,
"isStream": False,
"isSeekable": True,
"position": 0,
"thumbnail": track.image,
"isrc": track.isrc,
},
requester=self.bot.user,
)
for track in results
]
track = tracks[0]
return await self.get_recommendations(track=track, ctx=ctx)
class NodePool:
"""The base class for the node pool.
@ -1028,16 +930,15 @@ class NodePool:
password: str,
identifier: str,
secure: bool = False,
heartbeat: int = 120,
resume_key: Optional[str] = None,
resume_timeout: int = 60,
heartbeat: int = 30,
loop: Optional[asyncio.AbstractEventLoop] = None,
spotify_client_id: Optional[str] = None,
spotify_client_secret: Optional[str] = None,
session: Optional[aiohttp.ClientSession] = None,
apple_music: bool = False,
fallback: bool = False,
logger: Optional[logging.Logger] = None,
log_level: LogLevel = LogLevel.INFO,
log_handler: Optional[logging.Handler] = None,
) -> Node:
"""Creates a Node object to be then added into the node pool.
For Spotify searching capabilites, pass in valid Spotify API credentials.
@ -1056,15 +957,14 @@ class NodePool:
identifier=identifier,
secure=secure,
heartbeat=heartbeat,
resume_key=resume_key,
resume_timeout=resume_timeout,
loop=loop,
spotify_client_id=spotify_client_id,
session=session,
spotify_client_secret=spotify_client_secret,
apple_music=apple_music,
fallback=fallback,
logger=logger,
log_level=log_level,
log_handler=log_handler,
)
await node.connect()

View File

@ -203,13 +203,9 @@ class Queue(Iterable[Track]):
raise QueueEmpty("No items in the queue.")
if self._loop_mode == LoopMode.QUEUE:
# set current item to first track in queue if not set already
# otherwise exception will be raised
if not self._current_item or self._current_item not in self._queue:
if self._queue:
item = self._queue[0]
else:
raise QueueEmpty("No items in the queue.")
# recurse if the item isnt in the queue
if self._current_item not in self._queue:
self.get()
# set current item to first track in queue if not set already
if not self._current_item:

View File

@ -1,16 +1,13 @@
from __future__ import annotations
import asyncio
import logging
import re
import time
from base64 import b64encode
from typing import AsyncGenerator
from typing import Dict
from typing import List
from typing import Optional
from typing import Union
from urllib.parse import quote
import aiohttp
import orjson as json
@ -24,10 +21,8 @@ __all__ = ("Client",)
GRANT_URL = "https://accounts.spotify.com/api/token"
REQUEST_URL = "https://api.spotify.com/v1/{type}s/{id}"
# Keep this in sync with URLRegex.SPOTIFY_URL (enums.py). Accept intl locale segment,
# optional trailing slash, and query parameters.
SPOTIFY_URL_REGEX = re.compile(
r"https?://open\.spotify\.com/(?:intl-[a-zA-Z-]+/)?(?P<type>album|playlist|track|artist)/(?P<id>[a-zA-Z0-9]+)(?:/)?(?:\?.*)?$",
r"https?://open.spotify.com/(?P<type>album|playlist|track|artist)/(?P<id>[a-zA-Z0-9]+)",
)
@ -37,39 +32,29 @@ class Client:
for any Spotify URL you throw at it.
"""
def __init__(
self,
client_id: str,
client_secret: str,
*,
playlist_concurrency: int = 10,
playlist_page_limit: Optional[int] = None,
) -> None:
self._client_id = client_id
self._client_secret = client_secret
def __init__(self, client_id: str, client_secret: str) -> None:
self._client_id: str = client_id
self._client_secret: str = client_secret
# HTTP session will be injected by Node
self.session: Optional[aiohttp.ClientSession] = None
self.session: aiohttp.ClientSession = None # type: ignore
self._bearer_token: Optional[str] = None
self._expiry: float = 0.0
self._auth_token = b64encode(f"{self._client_id}:{self._client_secret}".encode())
self._grant_headers = {"Authorization": f"Basic {self._auth_token.decode()}"}
self._auth_token = b64encode(
f"{self._client_id}:{self._client_secret}".encode(),
)
self._grant_headers = {
"Authorization": f"Basic {self._auth_token.decode()}",
}
self._bearer_headers: Optional[Dict] = None
self._log = logging.getLogger(__name__)
# Performance tuning knobs
self._playlist_concurrency = max(1, playlist_concurrency)
self._playlist_page_limit = playlist_page_limit
async def _set_session(self, session: aiohttp.ClientSession) -> None:
self.session = session
async def _fetch_bearer_token(self) -> None:
_data = {"grant_type": "client_credentials"}
if not self.session:
raise SpotifyRequestException("HTTP session not initialized for Spotify client.")
resp = await self.session.post(GRANT_URL, data=_data, headers=self._grant_headers)
if resp.status != 200:
@ -78,8 +63,7 @@ class Client:
)
data: dict = await resp.json(loads=json.loads)
if self._log:
self._log.debug(f"Fetched Spotify bearer token successfully")
self._log.debug(f"Fetched Spotify bearer token successfully")
self._bearer_token = data["access_token"]
self._expiry = time.time() + (int(data["expires_in"]) - 10)
@ -100,8 +84,6 @@ class Client:
request_url = REQUEST_URL.format(type=spotify_type, id=spotify_id)
if not self.session:
raise SpotifyRequestException("HTTP session not initialized for Spotify client.")
resp = await self.session.get(request_url, headers=self._bearer_headers)
if resp.status != 200:
raise SpotifyRequestException(
@ -109,18 +91,15 @@ class Client:
)
data: dict = await resp.json(loads=json.loads)
if self._log:
self._log.debug(
f"Made request to Spotify API with status {resp.status} and response {data}",
)
self._log.debug(
f"Made request to Spotify API with status {resp.status} and response {data}",
)
if spotify_type == "track":
return Track(data)
elif spotify_type == "album":
return Album(data)
elif spotify_type == "artist":
if not self.session:
raise SpotifyRequestException("HTTP session not initialized for Spotify client.")
resp = await self.session.get(
f"{request_url}/top-tracks?market=US",
headers=self._bearer_headers,
@ -134,178 +113,37 @@ class Client:
tracks = track_data["tracks"]
return Artist(data, tracks)
else:
# For playlists we optionally use a reduced fields payload to shrink response sizes.
# NB: We cannot apply fields filter to initial request because original metadata is needed.
tracks = [
Track(track["track"])
for track in data["tracks"]["items"]
if track["track"] is not None
]
if not tracks:
if not len(tracks):
raise SpotifyRequestException(
"This playlist is empty and therefore cannot be queued.",
)
total_tracks = data["tracks"]["total"]
limit = data["tracks"]["limit"]
next_page_url = data["tracks"]["next"]
# Shortcircuit small playlists (single page)
if total_tracks <= limit:
return Playlist(data, tracks)
# Build remaining page URLs; Spotify supports offset-based pagination.
remaining_offsets = range(limit, total_tracks, limit)
page_urls: List[str] = []
fields_filter = (
"items(track(name,duration_ms,id,is_local,external_urls,external_ids,artists(name),album(images)))"
",next"
)
for idx, offset in enumerate(remaining_offsets):
if self._playlist_page_limit is not None and idx >= self._playlist_page_limit:
break
page_urls.append(
f"{request_url}/tracks?offset={offset}&limit={limit}&fields={quote(fields_filter)}",
)
if page_urls:
semaphore = asyncio.Semaphore(self._playlist_concurrency)
async def fetch_page(url: str) -> Optional[List[Track]]:
async with semaphore:
if not self.session:
raise SpotifyRequestException(
"HTTP session not initialized for Spotify client.",
)
resp = await self.session.get(url, headers=self._bearer_headers)
if resp.status != 200:
if self._log:
self._log.warning(
f"Page fetch failed {resp.status} {resp.reason} for {url}",
)
return None
page_json: dict = await resp.json(loads=json.loads)
return [
Track(item["track"])
for item in page_json.get("items", [])
if item.get("track") is not None
]
# Chunk gather in waves to avoid creating thousands of tasks at once
aggregated: List[Track] = []
wave_size = self._playlist_concurrency * 2
for i in range(0, len(page_urls), wave_size):
wave = page_urls[i : i + wave_size]
results = await asyncio.gather(
*[fetch_page(url) for url in wave],
return_exceptions=False,
while next_page_url is not None:
resp = await self.session.get(next_page_url, headers=self._bearer_headers)
if resp.status != 200:
raise SpotifyRequestException(
f"Error while fetching results: {resp.status} {resp.reason}",
)
for result in results:
if result:
aggregated.extend(result)
tracks.extend(aggregated)
next_data: dict = await resp.json(loads=json.loads)
tracks += [
Track(track["track"])
for track in next_data["items"]
if track["track"] is not None
]
next_page_url = next_data["next"]
return Playlist(data, tracks)
async def iter_playlist_tracks(
self,
*,
query: str,
batch_size: int = 100,
) -> AsyncGenerator[List[Track], None]:
"""Stream playlist tracks in batches without waiting for full materialization.
Parameters
----------
query: str
Spotify playlist URL.
batch_size: int
Number of tracks yielded per batch (logical grouping after fetch). Does not alter API page size.
"""
if not self._bearer_token or time.time() >= self._expiry:
await self._fetch_bearer_token()
match = SPOTIFY_URL_REGEX.match(query)
if not match or match.group("type") != "playlist":
raise InvalidSpotifyURL("Provided query is not a valid Spotify playlist URL.")
playlist_id = match.group("id")
request_url = REQUEST_URL.format(type="playlist", id=playlist_id)
if not self.session:
raise SpotifyRequestException("HTTP session not initialized for Spotify client.")
resp = await self.session.get(request_url, headers=self._bearer_headers)
if resp.status != 200:
raise SpotifyRequestException(
f"Error while fetching results: {resp.status} {resp.reason}",
)
data: dict = await resp.json(loads=json.loads)
# Yield first page immediately
first_page_tracks = [
Track(item["track"])
for item in data["tracks"]["items"]
if item.get("track") is not None
]
# Batch yield
for i in range(0, len(first_page_tracks), batch_size):
yield first_page_tracks[i : i + batch_size]
total = data["tracks"]["total"]
limit = data["tracks"]["limit"]
remaining_offsets = range(limit, total, limit)
fields_filter = (
"items(track(name,duration_ms,id,is_local,external_urls,external_ids,artists(name),album(images)))"
",next"
)
semaphore = asyncio.Semaphore(self._playlist_concurrency)
async def fetch(offset: int) -> List[Track]:
url = (
f"{request_url}/tracks?offset={offset}&limit={limit}&fields={quote(fields_filter)}"
)
async with semaphore:
if not self.session:
raise SpotifyRequestException(
"HTTP session not initialized for Spotify client.",
)
r = await self.session.get(url, headers=self._bearer_headers)
if r.status != 200:
if self._log:
self._log.warning(
f"Skipping page offset={offset} due to {r.status} {r.reason}",
)
return []
pj: dict = await r.json(loads=json.loads)
return [
Track(item["track"])
for item in pj.get("items", [])
if item.get("track") is not None
]
# Fetch pages in rolling waves; yield promptly as soon as a wave completes.
wave_size = self._playlist_concurrency * 2
for i, offset in enumerate(remaining_offsets):
# Build wave
if i % wave_size == 0:
wave_offsets = list(
o for o in remaining_offsets if o >= offset and o < offset + wave_size
)
results = await asyncio.gather(*[fetch(o) for o in wave_offsets])
for page_tracks in results:
if not page_tracks:
continue
for j in range(0, len(page_tracks), batch_size):
yield page_tracks[j : j + batch_size]
# Skip ahead in iterator by adjusting enumerate drive (consume extras)
# Fast-forward the generator manually
for _ in range(len(wave_offsets) - 1):
try:
next(remaining_offsets) # type: ignore
except StopIteration:
break
async def get_recommendations(self, *, query: str) -> List[Track]:
if not self._bearer_token or time.time() >= self._expiry:
await self._fetch_bearer_token()
@ -327,8 +165,6 @@ class Client:
id=f"?seed_tracks={spotify_id}",
)
if not self.session:
raise SpotifyRequestException("HTTP session not initialized for Spotify client.")
resp = await self.session.get(request_url, headers=self._bearer_headers)
if resp.status != 200:
raise SpotifyRequestException(
@ -339,22 +175,3 @@ class Client:
tracks = [Track(track) for track in data["tracks"]]
return tracks
async def track_search(self, *, query: str) -> List[Track]:
if not self._bearer_token or time.time() >= self._expiry:
await self._fetch_bearer_token()
request_url = f"https://api.spotify.com/v1/search?q={quote(query)}&type=track"
if not self.session:
raise SpotifyRequestException("HTTP session not initialized for Spotify client.")
resp = await self.session.get(request_url, headers=self._bearer_headers)
if resp.status != 200:
raise SpotifyRequestException(
f"Error while fetching results: {resp.status} {resp.reason}",
)
data: dict = await resp.json(loads=json.loads)
tracks = [Track(track) for track in data["tracks"]["items"]]
return tracks

View File

@ -4,7 +4,7 @@ import re
import setuptools
version = ""
requirements = ["aiohttp>=3.7.4,<4", "orjson", "websockets"]
requirements = ["aiohttp>=3.7.4,<4", "orjson"]
with open("pomice/__init__.py") as f:
version = re.search(
r'^__version__\s*=\s*[\'"]([^\'"]*)[\'"]',