From 045ec9ff668cf6a40bd02d74a2f979ed76c8e4dd Mon Sep 17 00:00:00 2001 From: insert Date: Sun, 1 Dec 2024 17:06:09 -0500 Subject: [PATCH] use sql for queue --- newbot.py | 144 ++++++++++++++++++++++++++++++------------------------ 1 file changed, 79 insertions(+), 65 deletions(-) diff --git a/newbot.py b/newbot.py index 500f0ea..dda24bc 100644 --- a/newbot.py +++ b/newbot.py @@ -9,27 +9,21 @@ from disnake.ext import commands, tasks from disnake import TextInputStyle import asyncio from dotenv import load_dotenv -import os +import os, sys import datetime import logging import math from tempfile import TemporaryDirectory from pathlib import Path -#https://stackoverflow.com/a/11706463 -def randshuffle(x, start=0, stop=None, seed=10): - if stop is None: - stop = len(x) +import sqlite3 - loop = 0 - for i in reversed(range(start + 1, stop)): - random.seed(seed+loop) - j = random.randint(start, i) - x[i], x[j] = x[j], x[i] - loop = loop + 1 +if len(sys.argv) > 1: + if sys.argv[1] == "--clear-queue": + os.remove("queue.db") + print("Queue cleared!") load_dotenv() queue = [] -user_queue = [] skip_list = [] shuffle = True retries = 0 @@ -40,6 +34,29 @@ vid_dir = Path(temp_dir.name) download_dir = TemporaryDirectory() full_dl_dir = Path(download_dir.name) vid_details = {"title": "", "channel": ""} +con = sqlite3.connect("queue.db", check_same_thread=False) +cur = con.cursor() + +def countuser(usrid): + res = cur.execute(f"SELECT ROWID FROM queue WHERE " + ("hasplayed = false AND " if not (os.getenv("PERMANENT_MAX_QUEUE","FALSE") == "TRUE") else "") + f"user = {usrid}") + rowid = res.fetchall() + print(len(rowid)) + return len(rowid) + +def sqllen(): + res = cur.execute(f"SELECT ROWID FROM queue WHERE hasplayed = false") + rowid = res.fetchall() + return len(rowid) + +def propagate_queue(times): + res = cur.execute(f"SELECT ROWID,link FROM queue WHERE hasplayed = false " + ("ORDER BY RANDOM() " if shuffle else "") + f"LIMIT {times}") + for rowid,link in res.fetchall(): + if len(queue) > 2: + print(f"The queue is larger than two videos this WILL cause issues: {queue}") + queue.append(link) + cur.execute(f"UPDATE queue SET hasplayed = true WHERE ROWID='{rowid}'") + print("Updated a line!!") + con.commit() def video_check(info, *, incomplete): duration = info.get('duration') @@ -49,8 +66,7 @@ def video_check(info, *, incomplete): print(info.get("title")) if duration and duration >= ((float(os.getenv("MAX_MIN")) * 60) + 480): queue.pop(0) - if not (os.getenv("PERMANENT_MAX_QUEUE","FALSE") == "TRUE"): - user_queue.pop(0) + propagate_queue(1) skip_list.clear() if queue: ydl.download(queue[0]) @@ -60,7 +76,7 @@ ydl_opts = { 'match_filter': video_check, 'hls_prefer_native': True, 'extract_flat': 'discard_in_playlist', - 'format': 'bestvideo[height<=1080][ext=mp4]+bestaudio[abr<=256][ext=m4a]/best[ext=mp4]/best', + 'format': 'bestvideo[height<=1080][vcodec!*=av01][ext=mp4]+bestaudio[abr<=256][ext=m4a]/best[ext=mp4]/best', 'fragment_retries': 10, 'noplaylist': True, 'restrictfilenames': True, @@ -107,13 +123,18 @@ print(f"OBS Version: {ver.obs_version}") @bot.event async def on_ready(): + res = cur.execute("SELECT name FROM sqlite_master WHERE name='queue'") + if res.fetchone() is None: + cur.execute("CREATE TABLE queue(link, user, hasplayed)") + con.commit() obs.set_current_program_scene("waiting") obse.callback.register(on_media_input_playback_ended) videotimer.start() ensurewaiting.start() titlehandler.start() - if (not os.path.isfile(f"{vid_dir}/{vidcounter}.mp4")) and len(queue) > 1: + if (not os.path.isfile(f"{vid_dir}/{vidcounter}.mp4")) and sqllen() > 1: print("predefined queue!") + propagate_queue(2) loop = asyncio.get_running_loop() await loop.run_in_executor(None, cold_run) @@ -131,8 +152,7 @@ def download_video(index): global failures if retries % 2 == 1: queue.pop(index) - if not (os.getenv("PERMANENT_MAX_QUEUE","FALSE") == "TRUE"): - user_queue.pop(index) + propagate_queue(1) failures = failures + 1 retries = retries + 1 download_video(index) @@ -176,11 +196,10 @@ def on_media_input_playback_ended(data): global vidcounter try: queue.pop(0) - except IndexError: - print("Video ended but somehow it wasn't in the queue? Resetting") + propagate_queue(1) + except IndexError as e: + print(f"Video ended but somehow it wasn't in the queue? Resetting {e}") obs.set_current_program_scene("waiting") - if not (os.getenv("PERMANENT_MAX_QUEUE","FALSE") == "TRUE"): - user_queue.pop(0) skip_list.clear() if not queue: obs.set_current_program_scene("waiting") @@ -189,14 +208,6 @@ def on_media_input_playback_ended(data): return wait_for_next_video() print(queue) - print(user_queue) - if shuffle: - seed = time.time() - randshuffle(queue,1,None,seed) - if not (os.getenv("PERMANENT_MAX_QUEUE","FALSE") == "TRUE"): - randshuffle(user_queue,1,None,seed) - print(queue) - print(user_queue) #obs.set_current_program_scene("youtube2") os.remove(f"{vid_dir}/{vidcounter}.mp4") vidcounter = vidcounter + 1 @@ -233,10 +244,15 @@ async def stats(inter: disnake.AppCmdInter): message = message + f"Failures: {failures}\n" nowplaying = obs.get_input_settings("nowplaying") message = message + f"Title: {nowplaying.input_settings['text']}\n" + message = message + f"Link: <{queue[0]}>\n" playing = obs.get_media_input_status("player") message = message + f"Video Duration: {str(datetime.timedelta(seconds=(round(playing.media_cursor/1000))))}/{str(datetime.timedelta(seconds=(round(playing.media_duration/1000))))}\n" if inter.permissions.moderate_members and not (os.getenv("PERMANENT_MAX_QUEUE","FALSE") == "TRUE"): - message = message + f"Queued by <@{user_queue[0]}>" + res = cur.execute(f"SELECT user FROM queue WHERE hasplayed = true AND link = '{queue[0]}'") + res = res.fetchall() #We can't gaurentee the result so just show likely possibilites + message = message + f"Users who have queued this video: " + for [usrid] in list(set(res)): + message = message + f"<@{usrid}>, " await inter.edit_original_response(message) @@ -246,15 +262,20 @@ async def stats(inter: disnake.AppCmdInter): ) async def play(inter: disnake.AppCmdInter, link: str): await inter.response.defer(ephemeral=True) - if user_queue.count(inter.user.id) >= (int(os.getenv("MAX_QUEUE"))): + if countuser(inter.user.id) >= (int(os.getenv("MAX_QUEUE"))): await inter.edit_original_response(f"You have reached the queue limit of {os.getenv('MAX_QUEUE')}, " + ("try again after one of your videos has played." if not (os.getenv("PERMANENT_MAX_QUEUE","FALSE") == "TRUE") else "you may not queue videos for the rest of the session.")) return if urlparse(link).netloc == 'youtube.com' or urlparse(link).netloc == 'www.youtube.com' or urlparse(link).netloc == 'youtu.be': - queue.append(link) - user_queue.append(inter.user.id) + cur.execute(f"""INSERT INTO queue VALUES + ('{link}',{inter.user.id},false) + """) + con.commit() await inter.edit_original_response(f"added to queue!") - if (not os.path.isfile(f"{vid_dir}/{vidcounter}.mp4")) and len(queue) > 1: + if (not os.path.isfile(f"{vid_dir}/{vidcounter}.mp4")) and sqllen() > 1 and len(queue) == 0: loop = asyncio.get_running_loop() + queue.clear() #safety + propagate_queue(2) + print(queue) await loop.run_in_executor(None, cold_run) return else: @@ -290,16 +311,20 @@ async def getqueue(inter: disnake.AppCmdInter): await inter.edit_original_response("There are no items in queue") return message = f"Now Playing: <{queue[0]}>\n" + message = message + f"Up Next: <{queue[1]}>\n" message = message + f"Shuffle is currently " + ("off\n" if not shuffle else "on!\n") - for i in range(11): + res = cur.execute(f"SELECT link FROM queue WHERE hasplayed = false") + links = [x[0] for x in res.fetchall()] + links.insert(0,"dummy") + for i in range(10): if i == 0: continue try: - message = message + f"{i}. <{queue[i]}>\n" + message = message + f"{i}. <{links[i]}>\n" except IndexError: break - message = message + f"1 of {math.ceil((len(queue)-1)/10) if (len(queue)-1)/10 > 1 else 1}" - if math.ceil((len(queue)-1)/10) > 1: + message = message + f"1 of {math.ceil((len(links)-1)/10) if (len(links)-1)/10 > 1 else 1}" + if math.ceil((len(links)-1)/10) > 1: await inter.edit_original_response(message, components=[disnake.ui.Button(label=">>", style=disnake.ButtonStyle.primary, custom_id="Forward"),]) else: await inter.edit_original_response(message) @@ -314,16 +339,20 @@ async def button_listener(inter: disnake.MessageInteraction): page = page[-1].split(" of ") if inter.component.custom_id == "Forward": message = f"Now Playing: <{queue[0]}>\n" + message = message + f"Up Next: <{queue[1]}>\n" message = message + f"Shuffle is currently " + ("off\n" if not shuffle else "on!\n") + res = cur.execute(f"SELECT link FROM queue WHERE hasplayed = false") + links = [x[0] for x in res.fetchall()] + links.insert(0,"dummy") offset = int(int(page[0]) * 10) for i in range(11): if i == 0: continue try: - message = message + f"{int(i+offset)}. <{queue[int(i+offset)]}>\n" + message = message + f"{int(i+offset)}. <{links[int(i+offset)]}>\n" except IndexError: break - message = message + f"{int(page[0])+1} of {math.ceil((len(queue)-1)/10) if (len(queue)-1)/10 > 1 else 1}" + message = message + f"{int(page[0])+1} of {math.ceil((len(links)-1)/10) if (len(links)-1)/10 > 1 else 1}" if (int(page[0])+1) >= int(page[1]): await inter.response.edit_message(message, components=[disnake.ui.Button(label="<<", style=disnake.ButtonStyle.primary, custom_id="Backward"),]) else: @@ -332,16 +361,20 @@ async def button_listener(inter: disnake.MessageInteraction): return if inter.component.custom_id == "Backward": message = f"Now Playing: <{queue[0]}>\n" + message = message + f"Up Next: <{queue[1]}>\n" message = message + f"Shuffle is currently " + ("off\n" if not shuffle else "on!\n") + res = cur.execute(f"SELECT link FROM queue WHERE hasplayed = false") + links = [x[0] for x in res.fetchall()] + links.insert(0,"dummy") offset = int((int(page[0]) - 2) * 10) for i in range(11): if i == 0: continue try: - message = message + f"{int(i+offset)}. <{queue[int(i+offset)]}>\n" + message = message + f"{int(i+offset)}. <{links[int(i+offset)]}>\n" except IndexError: break - message = message + f"{int(page[0])-1} of {math.ceil((len(queue)-1)/10) if (len(queue)-1)/10 > 1 else 1}" + message = message + f"{int(page[0])-1} of {math.ceil((len(links)-1)/10) if (len(links)-1)/10 > 1 else 1}" if (int(page[0])-1) <= 1 and int(int(page[1]) == 1): await inter.response.edit_message(message) elif (int(page[0])-1) <= 1 and int(int(page[1]) > 1): @@ -378,8 +411,7 @@ async def skip(inter: disnake.AppCmdInter): loop = asyncio.get_running_loop() global vidcounter queue.pop(0) - if not (os.getenv("PERMANENT_MAX_QUEUE","FALSE") == "TRUE"): - user_queue.pop(0) + propagate_queue(1) skip_list.clear() if not queue: obs.set_current_program_scene("waiting") @@ -436,8 +468,7 @@ async def voteskip(inter: disnake.AppCmdInter): loop = asyncio.get_running_loop() global vidcounter queue.pop(0) - if not (os.getenv("PERMANENT_MAX_QUEUE","FALSE") == "TRUE"): - user_queue.pop(0) + propagate_queue(1) skip_list.clear() if not queue: obs.set_current_program_scene("waiting") @@ -463,22 +494,6 @@ async def voteskip(inter: disnake.AppCmdInter): await inter.edit_original_response(f"**{inter.user.display_name}** has voted to skip the video, {len(skip_list)}/{math.floor(len(vc)/2)}") -@bot.slash_command( - name="remove", - description="removes a video from the queue", - default_member_permissions=disnake.Permissions(8192), -) -async def remove(inter: disnake.AppCmdInter, toremove: int): - await inter.response.defer(ephemeral=True) - if toremove == 0 or toremove == 1: - await inter.edit_original_response("that is the currently playing video!", ephemeral=True) - return - else: - queue.pop(toremove) - if not (os.getenv("PERMANENT_MAX_QUEUE","FALSE") == "TRUE"): - user_queue.pop(toremove) - await inter.edit_original_response("removed!") - @tasks.loop(seconds=5.0) async def videotimer(): @@ -497,8 +512,7 @@ async def videotimer(): print("skip return") return queue.pop(0) - if not (os.getenv("PERMANENT_MAX_QUEUE","FALSE") == "TRUE"): - user_queue.pop(0) + propagate_queue(1) skip_list.clear() if not queue: obs.set_current_program_scene("waiting")