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 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()
|
||||||
)
|
)
|
||||||
|
|
Loading…
Reference in a new issue