diff options
| -rw-r--r-- | bot/exts/fun/latex/_latex_cog.py | 86 | ||||
| -rw-r--r-- | bot/exts/fun/latex/_renderer.py | 47 | ||||
| -rw-r--r-- | bot/exts/fun/latex/template.txt | 5 |
3 files changed, 63 insertions, 75 deletions
diff --git a/bot/exts/fun/latex/_latex_cog.py b/bot/exts/fun/latex/_latex_cog.py index 0cd4c981..72d48b2a 100644 --- a/bot/exts/fun/latex/_latex_cog.py +++ b/bot/exts/fun/latex/_latex_cog.py @@ -1,13 +1,14 @@ -import asyncio import hashlib import re -import sys +import string from pathlib import Path -from textwrap import dedent +from typing import BinaryIO, Optional import discord from discord.ext import commands +from bot.bot import Bot + FORMATTED_CODE_REGEX = re.compile( r"(?P<delim>(?P<block>```)|``?)" # code delimiter: 1-3 backticks; (?P=block) only matches if it's a block r"(?(block)(?:(?P<lang>[a-z]+)\n)?)" # if we're in a block, match optional language (only letters plus newline) @@ -18,35 +19,65 @@ FORMATTED_CODE_REGEX = re.compile( 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 / "cache" CACHE_DIRECTORY.mkdir(exist_ok=True) +TEMPLATE = string.Template((THIS_DIR / "template.txt").read_text()) def _prepare_input(text: str) -> str: - text = text.replace(r"\\", "$\n$") # matplotlib uses \n for newlines, not \\ - if match := FORMATTED_CODE_REGEX.match(text): return match.group("code") else: return text -def _format_err(text: str) -> str: - # prevent escaping the codeblock by inserting a zero-width-joiner - text = text.replace("`", "`\u200d") - return dedent( - f""" - ``` - {text} - ``` - """ - ).strip() +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: @@ -56,18 +87,17 @@ class Latex(commands.Cog): image_path = CACHE_DIRECTORY / f"{query_hash}.png" async with ctx.typing(): if not image_path.exists(): - proc = await asyncio.subprocess.create_subprocess_exec( - sys.executable, - "_renderer.py", - query, - image_path.relative_to(THIS_DIR), - cwd=THIS_DIR, - stderr=asyncio.subprocess.PIPE - ) - return_code = await proc.wait() - if return_code != 0: + 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() - err = (await proc.stderr.read()).decode() - raise commands.BadArgument(_format_err(err)) - + return await ctx.send(file=discord.File(image_path, "latex.png")) diff --git a/bot/exts/fun/latex/_renderer.py b/bot/exts/fun/latex/_renderer.py deleted file mode 100644 index fb72b94c..00000000 --- a/bot/exts/fun/latex/_renderer.py +++ /dev/null @@ -1,47 +0,0 @@ -import sys -from pathlib import Path -from typing import BinaryIO - -import matplotlib.pyplot as plt - -# configure fonts and colors for matplotlib -plt.rcParams.update( - { - "font.size": 16, - "mathtext.fontset": "cm", # Computer Modern font set - "mathtext.rm": "serif", - "figure.facecolor": "36393F", # matches Discord's dark mode background color - "text.color": "white", - } -) - - -def render(text: str, file_handle: BinaryIO) -> None: - """ - Saves rendered image in `file_handle`. - - In case the input is invalid latex, it prints the error to `stderr`. - """ - fig = plt.figure() - fig.text(0, 1, text, horizontalalignment="left", verticalalignment="top") - try: - plt.savefig(file_handle, bbox_inches="tight", dpi=600) - except ValueError as err: - # get rid of traceback, keeping just the latex error - sys.exit(err) - - -def main() -> None: - """ - Renders a latex query and saves the output in a specified file. - - Expects two command line arguments: the query and the path to the output file. - """ - query = sys.argv[1] - out_file_path = Path(sys.argv[2]) - with open(out_file_path, "wb") as out_file: - render(query, out_file) - - -if __name__ == "__main__": - main() diff --git a/bot/exts/fun/latex/template.txt b/bot/exts/fun/latex/template.txt new file mode 100644 index 00000000..a20cc279 --- /dev/null +++ b/bot/exts/fun/latex/template.txt @@ -0,0 +1,5 @@ +\documentclass{article} +\begin{document} + \pagenumbering{gobble} + $text +\end{document} |