Introduce resolution autocomplete
This commit is contained in:
parent
45e85cfba4
commit
f6d093d45d
1 changed files with 45 additions and 26 deletions
|
@ -6,6 +6,7 @@ import os
|
|||
import tempfile
|
||||
import time
|
||||
import copy
|
||||
import typing
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import discord
|
||||
|
@ -19,22 +20,41 @@ from selenium.webdriver.chrome.service import Service as ChromeService
|
|||
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):
|
||||
def __init__(self, bot: commands.Bot):
|
||||
self.bot = bot
|
||||
self.log = logging.getLogger("jimmy.cogs.screenshot")
|
||||
|
||||
self.chrome_options = ChromeOptions()
|
||||
self.chrome_options.add_argument("--headless")
|
||||
self.chrome_options.add_argument("--disable-extensions")
|
||||
self.chrome_options.add_argument("--incognito")
|
||||
self.chrome_options.add_argument("--remote-debugging-port=9222")
|
||||
def chrome_options(self) -> ChromeOptions:
|
||||
_chrome_options = ChromeOptions()
|
||||
_chrome_options.add_argument("--headless")
|
||||
_chrome_options.add_argument("--disable-extensions")
|
||||
_chrome_options.add_argument("--incognito")
|
||||
_chrome_options.add_argument("--remote-debugging-port=9222")
|
||||
if os.getuid() == 0:
|
||||
self.chrome_options.add_argument("--no-sandbox")
|
||||
self.chrome_options.add_argument("--disable-dev-shm-usage")
|
||||
self.chrome_options.add_argument("--disable-setuid-sandbox")
|
||||
_chrome_options.add_argument("--no-sandbox")
|
||||
_chrome_options.add_argument("--disable-dev-shm-usage")
|
||||
_chrome_options.add_argument("--disable-setuid-sandbox")
|
||||
self.log.warning("Running as root, disabling chrome sandbox.")
|
||||
self.chrome_options.add_argument(
|
||||
_chrome_options.add_argument(
|
||||
"--user-agent=Mozi11a/5.0 "
|
||||
"(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,
|
||||
"download_restrictions": 3,
|
||||
}
|
||||
self.chrome_options.add_experimental_option(
|
||||
_chrome_options.add_experimental_option(
|
||||
"prefs", prefs
|
||||
)
|
||||
return _chrome_options
|
||||
|
||||
def compress_png(self, input_file: io.BytesIO) -> io.BytesIO:
|
||||
img = Image.open(input_file)
|
||||
|
@ -79,7 +100,10 @@ class ScreenshotCog(commands.Cog):
|
|||
load_timeout: int = 10,
|
||||
render_timeout: int = None,
|
||||
eager: bool = None,
|
||||
resolution: str = "1920x1080",
|
||||
resolution: typing.Annotated[
|
||||
str,
|
||||
_RES_OPTION
|
||||
] = "1440p",
|
||||
use_proxy: bool = False
|
||||
):
|
||||
"""Screenshots a webpage."""
|
||||
|
@ -108,9 +132,12 @@ class ScreenshotCog(commands.Cog):
|
|||
|
||||
start_init = time.time()
|
||||
try:
|
||||
options = copy.copy(self.chrome_options)
|
||||
if use_proxy and (server := CONFIG["screenshot"].get("proxy")):
|
||||
options = self.chrome_options()
|
||||
if use_proxy is True and (server := CONFIG["screenshot"].get("proxy")):
|
||||
use_proxy = True
|
||||
options.add_argument("--proxy-server=" + server)
|
||||
else:
|
||||
use_proxy = False
|
||||
service = await asyncio.to_thread(ChromeService)
|
||||
driver: webdriver.Chrome = await asyncio.to_thread(
|
||||
webdriver.Chrome,
|
||||
|
@ -119,18 +146,7 @@ class ScreenshotCog(commands.Cog):
|
|||
)
|
||||
driver.set_page_load_timeout(load_timeout)
|
||||
if resolution:
|
||||
common = {
|
||||
"8k": "7680x4320",
|
||||
"4k": "3840x2160",
|
||||
"1440p": "2560x1440",
|
||||
"1080p": "1920x1080",
|
||||
"720p": "1280x720",
|
||||
"480p": "854x480",
|
||||
"360p": "640x360",
|
||||
"240p": "426x240",
|
||||
"144p": "256x144"
|
||||
}
|
||||
resolution = common.get(resolution.lower(), resolution)
|
||||
resolution = RESOLUTIONS.get(resolution.lower(), resolution)
|
||||
try:
|
||||
width, height = map(int, resolution.split("x"))
|
||||
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"Compress time: {seconds(start_compress, end_compress)}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(),
|
||||
timestamp=discord.utils.utcnow()
|
||||
)
|
||||
|
|
Loading…
Reference in a new issue