From 9d49fa0ae13900be5f14d237f92c43d11d85dfc9 Mon Sep 17 00:00:00 2001 From: Shakya Majumdar Date: Fri, 21 Jan 2022 20:21:33 +0530 Subject: rework directory structure --- bot/exts/fun/latex.py | 108 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 108 insertions(+) create mode 100644 bot/exts/fun/latex.py (limited to 'bot/exts/fun/latex.py') diff --git a/bot/exts/fun/latex.py b/bot/exts/fun/latex.py new file mode 100644 index 00000000..ac43e95b --- /dev/null +++ b/bot/exts/fun/latex.py @@ -0,0 +1,108 @@ +import hashlib +import re +import string +from pathlib import Path +from typing import BinaryIO, Optional + +import discord +from discord.ext import commands + +from bot.bot import Bot + +FORMATTED_CODE_REGEX = re.compile( + r"(?P(?P```)|``?)" # code delimiter: 1-3 backticks; (?P=block) only matches if it's a block + r"(?(block)(?:(?P[a-z]+)\n)?)" # if we're in a block, match optional language (only letters plus newline) + r"(?:[ \t]*\n)*" # any blank (empty or tabs/spaces only) lines before the code + r"(?P.*?)" # extract all code inside the markup + r"\s*" # any more whitespace before the end of the code markup + r"(?P=delim)", # match the exact same delimiter from the start again + re.DOTALL | re.IGNORECASE, # "." also matches newlines, case insensitive +) + +LATEX_API_URL = "https://rtex.probablyaweb.site/api/v2" +PASTEBIN_URL = "https://paste.pythondiscord.com" + +THIS_DIR = Path(__file__).parent +CACHE_DIRECTORY = THIS_DIR / "_latex_cache" +CACHE_DIRECTORY.mkdir(exist_ok=True) +TEMPLATE = string.Template(Path("bot/resources/fun/latex_template.txt").read_text()) + + +def _prepare_input(text: str) -> str: + if match := FORMATTED_CODE_REGEX.match(text): + return match.group("code") + else: + return text + + +class InvalidLatexError(Exception): + """Represents an error caused by invalid latex.""" + + def __init__(self, logs: str): + super().__init__(logs) + self.logs = logs + + +class Latex(commands.Cog): + """Renders latex.""" + + def __init__(self, bot: Bot): + self.bot = bot + + async def _generate_image(self, query: str, out_file: BinaryIO) -> None: + """Make an API request and save the generated image to cache.""" + payload = {"code": query, "format": "png"} + async with self.bot.http_session.post(LATEX_API_URL, data=payload, raise_for_status=True) as response: + response_json = await response.json() + if response_json["status"] != "success": + raise InvalidLatexError(logs=response_json["log"]) + async with self.bot.http_session.get( + f"{LATEX_API_URL}/{response_json['filename']}", + data=payload, + raise_for_status=True + ) as response: + out_file.write(await response.read()) + + async def _upload_to_pastebin(self, text: str) -> Optional[str]: + """Uploads `text` to the paste service, returning the url if successful.""" + try: + async with self.bot.http_session.post( + PASTEBIN_URL + "/documents", + data=text, + raise_for_status=True + ) as response: + response_json = await response.json() + if "key" in response_json: + return f"{PASTEBIN_URL}/{response_json['key']}.txt?noredirect" + except Exception: + # 400 (Bad Request) means there are too many characters + pass + + @commands.command() + @commands.max_concurrency(1, commands.BucketType.guild, wait=True) + async def latex(self, ctx: commands.Context, *, query: str) -> None: + """Renders the text in latex and sends the image.""" + query = _prepare_input(query) + query_hash = hashlib.md5(query.encode()).hexdigest() + image_path = CACHE_DIRECTORY / f"{query_hash}.png" + async with ctx.typing(): + if not image_path.exists(): + try: + with open(image_path, "wb") as out_file: + await self._generate_image(TEMPLATE.substitute(text=query), out_file) + except InvalidLatexError as err: + logs_paste_url = await self._upload_to_pastebin(err.logs) + embed = discord.Embed(title="Failed to render input.") + if logs_paste_url: + embed.description = f"[View Logs]({logs_paste_url})" + else: + embed.description = "Couldn't upload logs." + await ctx.send(embed=embed) + image_path.unlink() + return + await ctx.send(file=discord.File(image_path, "latex.png")) + + +def setup(bot: Bot) -> None: + """Load the Latex Cog.""" + bot.add_cog(Latex(bot)) -- cgit v1.2.3 From 7ae4792482f2a496a747f3b5892c623f71100f3f Mon Sep 17 00:00:00 2001 From: Shakya Majumdar Date: Wed, 26 Jan 2022 19:44:41 +0530 Subject: paste image on a white background --- bot/exts/fun/latex.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) (limited to 'bot/exts/fun/latex.py') diff --git a/bot/exts/fun/latex.py b/bot/exts/fun/latex.py index ac43e95b..80145843 100644 --- a/bot/exts/fun/latex.py +++ b/bot/exts/fun/latex.py @@ -1,10 +1,12 @@ import hashlib import re import string +from io import BytesIO from pathlib import Path from typing import BinaryIO, Optional import discord +from PIL import Image from discord.ext import commands from bot.bot import Bot @@ -27,6 +29,9 @@ CACHE_DIRECTORY = THIS_DIR / "_latex_cache" CACHE_DIRECTORY.mkdir(exist_ok=True) TEMPLATE = string.Template(Path("bot/resources/fun/latex_template.txt").read_text()) +BG_COLOR = (54, 57, 63, 255) +PAD = 10 + def _prepare_input(text: str) -> str: if match := FORMATTED_CODE_REGEX.match(text): @@ -35,6 +40,14 @@ def _prepare_input(text: str) -> str: return text +def _process_image(data: bytes, out_file: BinaryIO) -> None: + image = Image.open(BytesIO(data)).convert("RGBA") + width, height = image.size + background = Image.new("RGBA", (width + 2 * PAD, height + 2 * PAD), "WHITE") + background.paste(image, (PAD, PAD), image) + background.save(out_file) + + class InvalidLatexError(Exception): """Represents an error caused by invalid latex.""" @@ -58,10 +71,9 @@ class Latex(commands.Cog): raise InvalidLatexError(logs=response_json["log"]) async with self.bot.http_session.get( f"{LATEX_API_URL}/{response_json['filename']}", - data=payload, raise_for_status=True ) as response: - out_file.write(await response.read()) + _process_image(await response.read(), out_file) async def _upload_to_pastebin(self, text: str) -> Optional[str]: """Uploads `text` to the paste service, returning the url if successful.""" -- cgit v1.2.3 From 365fe989ee0b659b211ffc503f7ab3ca9641c7c9 Mon Sep 17 00:00:00 2001 From: Shakya Majumdar Date: Tue, 1 Feb 2022 10:34:31 +0530 Subject: add comments and docstrings, remove an unused variable --- bot/exts/fun/latex.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) (limited to 'bot/exts/fun/latex.py') diff --git a/bot/exts/fun/latex.py b/bot/exts/fun/latex.py index 80145843..1dc9c2e5 100644 --- a/bot/exts/fun/latex.py +++ b/bot/exts/fun/latex.py @@ -29,11 +29,11 @@ CACHE_DIRECTORY = THIS_DIR / "_latex_cache" CACHE_DIRECTORY.mkdir(exist_ok=True) TEMPLATE = string.Template(Path("bot/resources/fun/latex_template.txt").read_text()) -BG_COLOR = (54, 57, 63, 255) PAD = 10 def _prepare_input(text: str) -> str: + """Extract latex from a codeblock, if it is in one.""" if match := FORMATTED_CODE_REGEX.match(text): return match.group("code") else: @@ -41,9 +41,14 @@ def _prepare_input(text: str) -> str: def _process_image(data: bytes, out_file: BinaryIO) -> None: + """Read `data` as an image file, and paste it on a white background.""" image = Image.open(BytesIO(data)).convert("RGBA") width, height = image.size background = Image.new("RGBA", (width + 2 * PAD, height + 2 * PAD), "WHITE") + + # paste the image on the background, using the same image as the mask + # when an RGBA image is passed as the mask, its alpha band is used. + # this has the effect of skipping pasting the pixels where the image is transparent. background.paste(image, (PAD, PAD), image) background.save(out_file) @@ -95,6 +100,8 @@ class Latex(commands.Cog): async def latex(self, ctx: commands.Context, *, query: str) -> None: """Renders the text in latex and sends the image.""" query = _prepare_input(query) + + # the hash of the query is used as the filename in the cache. query_hash = hashlib.md5(query.encode()).hexdigest() image_path = CACHE_DIRECTORY / f"{query_hash}.png" async with ctx.typing(): -- cgit v1.2.3