klemek/meme-otron
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

checking image file size beforehand

Browse Source
This commit is contained in:
klemek
2020-04-30 08:55:12 +02:00
parent 671f6fd595
commit 59a8530cbe
2 changed files with 33 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
discord_bot/__main__.py
+22 -24
View File
@@ -121,7 +121,8 @@ async def on_message(message: discord.Message):
if len(message.attachments) > 0: if len(message.attachments) > 0:
input_data = await message.attachments[0].read() input_data = await message.attachments[0].read()
img, errors = meme_otron.compute(*args, left_wmark_text=left_wmark_text, input_data=input_data) img, errors = meme_otron.compute(*args, left_wmark_text=left_wmark_text,
input_data=input_data, max_file_size=8 * 1024 * 1024)
if len(errors) > 0: if len(errors) > 0:
response = ":warning:" response = ":warning:"
for err in errors: for err in errors:
@@ -135,29 +136,26 @@ async def on_message(message: discord.Message):
else: else:
with tempfile.NamedTemporaryFile(delete=False) as output: with tempfile.NamedTemporaryFile(delete=False) as output:
img.save(output, format="JPEG") img.save(output, format="JPEG")
if os.stat(output.name).st_size > 8 * 1024 * 1024: # 8MB response = None
await message.channel.send(":warning:\nOutput image is too big to be sent by discord") meme_id = utils.sanitize_input(args[0])
else: if len(args) == 1 and meme_id not in ["image", "text"]:
response = None meme = meme_db.get_meme(meme_id)
meme_id = utils.sanitize_input(args[0]) response = f"Template `{meme.id}`:"
if len(args) == 1 and meme_id not in ["image", "text"]: if len(meme.aliases) > 0:
meme = meme_db.get_meme(meme_id) response += f"\n- Aliases: `{'`, `'.join(meme.aliases)}`"
response = f"Template `{meme.id}`:" if meme.info is not None:
if len(meme.aliases) > 0: response += f"\n- More info: <{meme.info}>"
response += f"\n- Aliases: `{'`, `'.join(meme.aliases)}`" response += f"\n- Use:" \
if meme.info is not None: f"\n```{meme.id} \"" + \
response += f"\n- More info: <{meme.info}>" "\" \"".join([f"text {i + 1}" for i in range(meme.texts_len)]) + \
response += f"\n- Use:" \ "\"```"
f"\n```{meme.id} \"" + \ elif not is_direct:
"\" \"".join([f"text {i + 1}" for i in range(meme.texts_len)]) + \ response = f"A meme by {message.author.mention}:"
"\"```" if message_id not in SENT:
elif not is_direct: SENT[message_id] = []
response = f"A meme by {message.author.mention}:" response = await message.channel.send(response,
if message_id not in SENT: file=discord.File(filename="meme.jpg", fp=output.name))
SENT[message_id] = [] SENT[message_id] += [response]
response = await message.channel.send(response,
file=discord.File(filename="meme.jpg", fp=output.name))
SENT[message_id] += [response]
try: try:
os.remove(output.name) os.remove(output.name)
except PermissionError: except PermissionError:
meme_otron/meme_otron.py
+11 -2
View File
@@ -1,6 +1,7 @@
from typing import Optional, Tuple, List from typing import Optional, Tuple, List
import re import re
from PIL import Image from PIL import Image
from io import BytesIO
from .types import Text, Pos from .types import Text, Pos
from . import img_factory from . import img_factory
@@ -31,6 +32,7 @@ simple_text.y_range = [0.2, 0.8]
def compute(*args: str, input_data: Optional[bytes] = None, def compute(*args: str, input_data: Optional[bytes] = None,
wmark: bool = True, left_wmark_text: Optional[str] = None, wmark: bool = True, left_wmark_text: Optional[str] = None,
max_file_size: Optional[int] = None,
debug: bool = False) -> Tuple[Optional[Image.Image], List[str]]: debug: bool = False) -> Tuple[Optional[Image.Image], List[str]]:
if len(args) < 1: if len(args) < 1:
return None, ['Not enough arguments'] return None, ['Not enough arguments']
@@ -39,7 +41,7 @@ def compute(*args: str, input_data: Optional[bytes] = None,
images = [] images = []
errors = [] errors = []
for part in parts: for part in parts:
img, err = compute_part(*part, input_data=input_data, debug=debug) img, err = compute_part(*part, input_data=input_data, max_file_size=max_file_size, debug=debug)
if img is not None: if img is not None:
images += [img] images += [img]
else: else:
@@ -56,10 +58,17 @@ def compute(*args: str, input_data: Optional[bytes] = None,
watermarks += [left_wmark.variant(left_wmark_text)] watermarks += [left_wmark.variant(left_wmark_text)]
output_image = img_factory.apply_texts(output_image, watermarks, debug=debug) output_image = img_factory.apply_texts(output_image, watermarks, debug=debug)
if max_file_size is not None:
img_file = BytesIO()
output_image.save(img_file, 'jpg')
if img_file.tell() > max_file_size:
return None, ['Output image too big']
return output_image, errors return output_image, errors
def compute_part(*args: str, input_data: Optional[bytes] = None, def compute_part(*args: str, input_data: Optional[bytes] = None,
max_file_size: Optional[int] = None,
debug: bool = False) -> Tuple[Optional[Image.Image], Optional[str]]: debug: bool = False) -> Tuple[Optional[Image.Image], Optional[str]]:
meme_id = utils.sanitize_input(args[0]) meme_id = utils.sanitize_input(args[0])
@@ -73,7 +82,7 @@ def compute_part(*args: str, input_data: Optional[bytes] = None,
if len(args) <= 1: if len(args) <= 1:
return None, 'Image: received no input data nor URL' return None, 'Image: received no input data nor URL'
else: else:
input_data, err = utils.read_web(args[1]) input_data, err = utils.read_web(args[1], max_file_size=max_file_size)
if input_data is None: if input_data is None:
return None, 'Image: ' + err return None, 'Image: ' + err
img = img_factory.build_image_only(input_data) img = img_factory.build_image_only(input_data)