Add some more error handling to the screenshot command
This commit is contained in:
parent
97fcfd904b
commit
6fcd23ee35
4 changed files with 68 additions and 30 deletions
|
|
@ -4,8 +4,11 @@
|
|||
<option name="autoReloadType" value="SELECTIVE" />
|
||||
</component>
|
||||
<component name="ChangeListManager">
|
||||
<list default="true" id="aa6d62a8-d64d-4a60-a85f-8d9fa52b6b49" name="Changes" comment="Beta 1 * Screenshot (unpolished, functional) * Dig * Traceroute * Whois * Ping * yt-dl * Docker file">
|
||||
<list default="true" id="aa6d62a8-d64d-4a60-a85f-8d9fa52b6b49" name="Changes" comment="Properly exit the webdriver">
|
||||
<change afterPath="$PROJECT_DIR$/conf.py" afterDir="false" />
|
||||
<change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
|
||||
<change beforePath="$PROJECT_DIR$/cogs/screenshot.py" beforeDir="false" afterPath="$PROJECT_DIR$/cogs/screenshot.py" afterDir="false" />
|
||||
<change beforePath="$PROJECT_DIR$/main.py" beforeDir="false" afterPath="$PROJECT_DIR$/main.py" afterDir="false" />
|
||||
</list>
|
||||
<option name="SHOW_DIALOG" value="false" />
|
||||
<option name="HIGHLIGHT_CONFLICTS" value="true" />
|
||||
|
|
@ -68,6 +71,8 @@
|
|||
<workItem from="1704140889356" duration="166000" />
|
||||
<workItem from="1704141061129" duration="2495000" />
|
||||
<workItem from="1704224862150" duration="11897000" />
|
||||
<workItem from="1704296563227" duration="780000" />
|
||||
<workItem from="1704317484555" duration="1407000" />
|
||||
</task>
|
||||
<task id="LOCAL-00001" summary="Update gitignore">
|
||||
<option name="closed" value="true" />
|
||||
|
|
@ -93,7 +98,15 @@
|
|||
<option name="project" value="LOCAL" />
|
||||
<updated>1704236387023</updated>
|
||||
</task>
|
||||
<option name="localTasksCounter" value="4" />
|
||||
<task id="LOCAL-00004" summary="Update the docker image">
|
||||
<option name="closed" value="true" />
|
||||
<created>1704296742685</created>
|
||||
<option name="number" value="00004" />
|
||||
<option name="presentableId" value="LOCAL-00004" />
|
||||
<option name="project" value="LOCAL" />
|
||||
<updated>1704296742686</updated>
|
||||
</task>
|
||||
<option name="localTasksCounter" value="5" />
|
||||
<servers />
|
||||
</component>
|
||||
<component name="TypeScriptGeneratedFilesManager">
|
||||
|
|
@ -104,6 +117,7 @@
|
|||
<MESSAGE value="Update gitignore" />
|
||||
<MESSAGE value="Beta 1 * Screenshot (unpolished, functional) * Dig * Traceroute * Whois * Ping * yt-dl" />
|
||||
<MESSAGE value="Beta 1 * Screenshot (unpolished, functional) * Dig * Traceroute * Whois * Ping * yt-dl * Docker file" />
|
||||
<option name="LAST_COMMIT_MESSAGE" value="Beta 1 * Screenshot (unpolished, functional) * Dig * Traceroute * Whois * Ping * yt-dl * Docker file" />
|
||||
<MESSAGE value="Update the docker image" />
|
||||
<option name="LAST_COMMIT_MESSAGE" value="Update the docker image" />
|
||||
</component>
|
||||
</project>
|
||||
|
|
@ -18,15 +18,15 @@ from selenium.webdriver.chrome.service import Service as ChromeService
|
|||
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-dev-shm-usage")
|
||||
# self.chrome_options.add_argument("--disable-gpu")
|
||||
self.chrome_options.add_argument("--disable-extensions")
|
||||
self.chrome_options.add_argument("--incognito")
|
||||
if os.getuid() == 0:
|
||||
self.chrome_options.add_argument("--no-sandbox")
|
||||
self.log.warning("Running as root, disabling chrome sandbox.")
|
||||
|
||||
prefs = {
|
||||
"download.open_pdf_in_system_reader": False,
|
||||
|
|
@ -39,8 +39,6 @@ class ScreenshotCog(commands.Cog):
|
|||
"prefs", prefs
|
||||
)
|
||||
|
||||
self.log = logging.getLogger("jimmy.cogs.screenshot")
|
||||
|
||||
def compress_png(self, input_file: io.BytesIO) -> io.BytesIO:
|
||||
img = Image.open(input_file)
|
||||
img = img.convert("RGB")
|
||||
|
|
@ -83,32 +81,51 @@ class ScreenshotCog(commands.Cog):
|
|||
render_timeout = 30 if eager else 10
|
||||
if not url.startswith("http"):
|
||||
url = "https://" + url
|
||||
self.log.debug(
|
||||
"User %s (%s) is attempting to screenshot %r with load timeout %d, render timeout %d, %s loading, and "
|
||||
"a %s resolution.",
|
||||
ctx.author,
|
||||
ctx.author.id,
|
||||
url,
|
||||
load_timeout,
|
||||
render_timeout,
|
||||
"eager" if eager else "lazy",
|
||||
resolution
|
||||
)
|
||||
parsed = urlparse(url)
|
||||
await ctx.respond("Initialising...")
|
||||
|
||||
start_init = time.time()
|
||||
service = await asyncio.to_thread(ChromeService)
|
||||
driver: webdriver.Chrome = await asyncio.to_thread(
|
||||
webdriver.Chrome,
|
||||
service=service,
|
||||
options=self.chrome_options
|
||||
)
|
||||
driver.set_page_load_timeout(load_timeout)
|
||||
if resolution:
|
||||
try:
|
||||
width, height = map(int, resolution.split("x"))
|
||||
driver.set_window_size(width, height)
|
||||
if height > 4320 or width > 7680:
|
||||
return await ctx.respond("Invalid resolution. Max resolution is 7680x4320 (8K).")
|
||||
except ValueError:
|
||||
return await ctx.respond("Invalid resolution. please provide width x height, e.g. 1920x1080")
|
||||
if eager:
|
||||
driver.implicitly_wait(render_timeout)
|
||||
try:
|
||||
service = await asyncio.to_thread(ChromeService)
|
||||
driver: webdriver.Chrome = await asyncio.to_thread(
|
||||
webdriver.Chrome,
|
||||
service=service,
|
||||
options=self.chrome_options
|
||||
)
|
||||
driver.set_page_load_timeout(load_timeout)
|
||||
if resolution:
|
||||
try:
|
||||
width, height = map(int, resolution.split("x"))
|
||||
driver.set_window_size(width, height)
|
||||
if height > 4320 or width > 7680:
|
||||
return await ctx.respond("Invalid resolution. Max resolution is 7680x4320 (8K).")
|
||||
except ValueError:
|
||||
return await ctx.respond("Invalid resolution. please provide width x height, e.g. 1920x1080")
|
||||
if eager:
|
||||
driver.implicitly_wait(render_timeout)
|
||||
except Exception as e:
|
||||
await ctx.respond("Failed to initialise browser: " + str(e))
|
||||
raise
|
||||
end_init = time.time()
|
||||
|
||||
await ctx.edit(content=("Loading webpage..." if not eager else "Loading & screenshotting webpage..."))
|
||||
start_request = time.time()
|
||||
await asyncio.to_thread(driver.get, url)
|
||||
try:
|
||||
await asyncio.to_thread(driver.get, url)
|
||||
except Exception as e:
|
||||
await ctx.respond("Failed to get the webpage: " + str(e))
|
||||
raise
|
||||
end_request = time.time()
|
||||
|
||||
if not eager:
|
||||
|
|
@ -140,6 +157,7 @@ class ScreenshotCog(commands.Cog):
|
|||
|
||||
await ctx.edit(content="Cleaning up...")
|
||||
start_cleanup = time.time()
|
||||
await asyncio.to_thread(driver.close)
|
||||
await asyncio.to_thread(driver.quit)
|
||||
end_cleanup = time.time()
|
||||
|
||||
|
|
|
|||
10
conf.py
Normal file
10
conf.py
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
import toml
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
CONFIG = toml.load('config.toml')
|
||||
except FileNotFoundError:
|
||||
cwd = Path.cwd()
|
||||
logging.getLogger("jimmy.autoconf").critical("Unable to locate config.toml in %s.", cwd, exc_info=True)
|
||||
raise
|
||||
6
main.py
6
main.py
|
|
@ -11,11 +11,7 @@ from rich.logging import RichHandler
|
|||
|
||||
log = logging.getLogger("jimmy")
|
||||
|
||||
try:
|
||||
CONFIG = toml.load('config.toml')
|
||||
except FileNotFoundError:
|
||||
log.critical("Unable to locate config.toml.", exc_info=True)
|
||||
raise
|
||||
|
||||
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s %(levelname)s %(name)s %(message)s",
|
||||
|
|
|
|||
Loading…
Reference in a new issue