nex/college-bot-v2
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Activity Actions

Introduce resolution autocomplete

Browse source
This commit is contained in:
Nexus 2024-04-02 16:03:14 +01:00
parent 45e85cfba4
commit f6d093d45d
Signed by: nex
GPG key ID: 0FA334385D0B689F
1 changed files with 45 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

71
src/cogs/screenshot.py
View file

@ -6,6 +6,7 @@ import os
import tempfile import tempfile
import time import time
import copy import copy
import typing
from urllib.parse import urlparse from urllib.parse import urlparse
import discord import discord
@ -19,22 +20,41 @@ from selenium.webdriver.chrome.service import Service as ChromeService
from conf import CONFIG from conf import CONFIG
RESOLUTIONS = {
"8k": "7680x4320",
"4k": "3840x2160",
"1440p": "2560x1440",
"1080p": "1920x1080",
"720p": "1280x720",
"480p": "854x480",
"360p": "640x360",
"240p": "426x240",
"144p": "256x144"
}
_RES_OPTION = discord.Option(
name="resolution",
description="The resolution of the browser, can be WxH, or a preset (e.g. 1080p).",
autocomplete=discord.utils.basic_autocomplete(tuple(RESOLUTIONS.keys()))
)
class ScreenshotCog(commands.Cog): class ScreenshotCog(commands.Cog):
def __init__(self, bot: commands.Bot): def __init__(self, bot: commands.Bot):
self.bot = bot self.bot = bot
self.log = logging.getLogger("jimmy.cogs.screenshot") self.log = logging.getLogger("jimmy.cogs.screenshot")
self.chrome_options = ChromeOptions() def chrome_options(self) -> ChromeOptions:
self.chrome_options.add_argument("--headless") _chrome_options = ChromeOptions()
self.chrome_options.add_argument("--disable-extensions") _chrome_options.add_argument("--headless")
self.chrome_options.add_argument("--incognito") _chrome_options.add_argument("--disable-extensions")
self.chrome_options.add_argument("--remote-debugging-port=9222") _chrome_options.add_argument("--incognito")
_chrome_options.add_argument("--remote-debugging-port=9222")
if os.getuid() == 0: if os.getuid() == 0:
self.chrome_options.add_argument("--no-sandbox") _chrome_options.add_argument("--no-sandbox")
self.chrome_options.add_argument("--disable-dev-shm-usage") _chrome_options.add_argument("--disable-dev-shm-usage")
self.chrome_options.add_argument("--disable-setuid-sandbox") _chrome_options.add_argument("--disable-setuid-sandbox")
self.log.warning("Running as root, disabling chrome sandbox.") self.log.warning("Running as root, disabling chrome sandbox.")
self.chrome_options.add_argument( _chrome_options.add_argument(
"--user-agent=Mozi11a/5.0 " "--user-agent=Mozi11a/5.0 "
"(X11; Linux x86_64) AppleWebKit/537.36 (KHTML, Like Gecko) Chrome/121.9.6167.160 Safari/537.36" "(X11; Linux x86_64) AppleWebKit/537.36 (KHTML, Like Gecko) Chrome/121.9.6167.160 Safari/537.36"
) )
@ -44,9 +64,10 @@ class ScreenshotCog(commands.Cog):
"plugins.always_open_pdf_externally": False, "plugins.always_open_pdf_externally": False,
"download_restrictions": 3, "download_restrictions": 3,
} }
self.chrome_options.add_experimental_option( _chrome_options.add_experimental_option(
"prefs", prefs "prefs", prefs
) )
return _chrome_options
def compress_png(self, input_file: io.BytesIO) -> io.BytesIO: def compress_png(self, input_file: io.BytesIO) -> io.BytesIO:
img = Image.open(input_file) img = Image.open(input_file)
@ -79,7 +100,10 @@ class ScreenshotCog(commands.Cog):
load_timeout: int = 10, load_timeout: int = 10,
render_timeout: int = None, render_timeout: int = None,
eager: bool = None, eager: bool = None,
resolution: str = "1920x1080", resolution: typing.Annotated[
str,
_RES_OPTION
] = "1440p",
use_proxy: bool = False use_proxy: bool = False
): ):
"""Screenshots a webpage.""" """Screenshots a webpage."""
@ -108,9 +132,12 @@ class ScreenshotCog(commands.Cog):
start_init = time.time() start_init = time.time()
try: try:
options = copy.copy(self.chrome_options) options = self.chrome_options()
if use_proxy and (server := CONFIG["screenshot"].get("proxy")): if use_proxy is True and (server := CONFIG["screenshot"].get("proxy")):
use_proxy = True
options.add_argument("--proxy-server=" + server) options.add_argument("--proxy-server=" + server)
else:
use_proxy = False
service = await asyncio.to_thread(ChromeService) service = await asyncio.to_thread(ChromeService)
driver: webdriver.Chrome = await asyncio.to_thread( driver: webdriver.Chrome = await asyncio.to_thread(
webdriver.Chrome, webdriver.Chrome,
@ -119,18 +146,7 @@ class ScreenshotCog(commands.Cog):
) )
driver.set_page_load_timeout(load_timeout) driver.set_page_load_timeout(load_timeout)
if resolution: if resolution:
common = { resolution = RESOLUTIONS.get(resolution.lower(), resolution)
"8k": "7680x4320",
"4k": "3840x2160",
"1440p": "2560x1440",
"1080p": "1920x1080",
"720p": "1280x720",
"480p": "854x480",
"360p": "640x360",
"240p": "426x240",
"144p": "256x144"
}
resolution = common.get(resolution.lower(), resolution)
try: try:
width, height = map(int, resolution.split("x")) width, height = map(int, resolution.split("x"))
driver.set_window_size(width, height) driver.set_window_size(width, height)
@ -207,7 +223,10 @@ class ScreenshotCog(commands.Cog):
f"Save time: {seconds(start_save, end_save)}s\n" f"Save time: {seconds(start_save, end_save)}s\n"
f"Compress time: {seconds(start_compress, end_compress)}s\n" f"Compress time: {seconds(start_compress, end_compress)}s\n"
f"Cleanup time: {seconds(start_cleanup, end_cleanup)}s\n" f"Cleanup time: {seconds(start_cleanup, end_cleanup)}s\n"
f"Screenshot size: {screenshot_size_mb}MB\n", f"Total time: {seconds(start_init, end_cleanup)}s\n"
f"Screenshot size: {screenshot_size_mb}MB\n"
f"Resolution: {resolution}"
f"Used proxy: {use_proxy}",
colour=discord.Colour.dark_theme(), colour=discord.Colour.dark_theme(),
timestamp=discord.utils.utcnow() timestamp=discord.utils.utcnow()
) )