diff --git a/ubot/custom.py b/ubot/custom.py index b536a42..d70e816 100644 --- a/ubot/custom.py +++ b/ubot/custom.py @@ -42,7 +42,7 @@ class ExtendedNewMessage(NewMessage.Event): return (default, self) if return_msg else default - async def get_image(self, event=None, with_reply=True, force_reply=False): + async def get_image(self, event: NewMessage.Event = None, with_reply=True, force_reply=False, size_limit: int = 0): event = event or self if event and event.media and not force_reply: @@ -56,6 +56,9 @@ class ExtendedNewMessage(NewMessage.Event): if event.gif or event.video or event.audio or event.voice: return + if size_limit and event.file.size > size_limit: + return + return event.media.document if with_reply and event.is_reply: diff --git a/ubot/modules/deepfry.py b/ubot/modules/deepfry.py index 2effdce..dddd48e 100644 --- a/ubot/modules/deepfry.py +++ b/ubot/modules/deepfry.py @@ -43,7 +43,7 @@ async def deepfryer(event): except ValueError: frycount = 1 - data = await event.get_image() + data = await event.get_image(1024 * 3072) if not data: await event.reply("Reply to an image or sticker or caption an image to deep fry it!") diff --git a/ubot/modules/evaluation.py b/ubot/modules/evaluation.py index 92fda6e..875742d 100644 --- a/ubot/modules/evaluation.py +++ b/ubot/modules/evaluation.py @@ -120,7 +120,7 @@ def flipstickersync(sticker_webp_io): @ldr.add("stickimg", help="Converts images to sticker-sized PNG files.") async def createsticker(event): - data = await event.get_image() + data = await event.get_image(size_limit=1024 * 3072) if not data: await event.reply("Reply to or caption an image to make it sticker-sized!") @@ -133,7 +133,7 @@ async def createsticker(event): @ldr.add("emojiimg", help="Converts images to emoji-sized PNG files.") async def createemoji(event): - data = await event.get_image() + data = await event.get_image(size_limit=1024 * 3072) if not data: await event.reply("Reply to or caption an image to make it emoji-sized!")