diff --git a/jarvis/cogs/admin.py b/jarvis/cogs/admin.py index cba225b..55f1c03 100644 --- a/jarvis/cogs/admin.py +++ b/jarvis/cogs/admin.py @@ -3,6 +3,12 @@ from discord.ext import commands class AdminCog(commands.Cog): + """ + Guild admin functions + + Used to manage guilds + """ + def __init__(self, bot): self.bot = bot diff --git a/jarvis/cogs/image.py b/jarvis/cogs/image.py new file mode 100644 index 0000000..b8bac1e --- /dev/null +++ b/jarvis/cogs/image.py @@ -0,0 +1,74 @@ +import jarvis +from discord import File +import cv2 +import aiohttp +import re +import numpy as np +from jarvis.utils import convert_bytesize, unconvert_bytesize +from discord.ext import commands +from io import BytesIO + + +class ImageCog(commands.Cog): + """ + Image processing functions for J.A.R.V.I.S. + + May be categorized under util later + """ + + def __init__(self, bot): + self.bot = bot + self._session = aiohttp.ClientSession() + self.tgt_match = re.compile( + r"([0-9]*\.?[0-9]*?) ?([KMGTP]?B)", re.IGNORECASE + ) + + @commands.command(name="resize", help="Resize an image") + async def _resize(self, ctx, target: str, url: str = None): + if not target: + await ctx.send("Missing target size, i.e. 200KB.") + return + + tgt = self.tgt_match.match(target) + if not tgt: + await ctx.send( + "Invalid target format ({}).".format(target) + + " Expected format like 200KB" + ) + return + + tgt_size = unconvert_bytesize(*tgt.groups()) + file = None + if ( + ctx.message.attachments is not None + and len(ctx.message.attachments) > 0 + ): + file = await ctx.message.attachments[0].read() + elif url is not None: + async with self._session.get(url) as resp: + if resp.status == 200: + file = await resp.read() + else: + ctx.send("Missing file as either attachment or URL.") + + size = len(file) + if size <= tgt_size: + await ctx.send("Image already meets target.") + return + + ratio = size / tgt_size + + buffer = np.frombuffer(file, dtype=np.uint8) + img = cv2.imdecode(buffer, flags=1) + + width = int(img.shape[1] * ratio) + height = int(img.shape[0] * ratio) + + new_img = cv2.resize(img, (width, height)) + buffer = cv2.imencode(".png", new_img)[1].tobytes() + bufio = BytesIO(buffer) + await ctx.send(file=File(bufio)) + + +def setup(bot): + bot.add_cog(ImageCog(bot)) diff --git a/jarvis/cogs/jokes.py b/jarvis/cogs/jokes.py index d05cfb9..ea79b15 100644 --- a/jarvis/cogs/jokes.py +++ b/jarvis/cogs/jokes.py @@ -10,6 +10,12 @@ from datetime import datetime class JokeCog(commands.Cog): + """ + Joke library for J.A.R.V.I.S. + + May adapt over time to create jokes using machine learning + """ + def __init__(self, bot): self.bot = bot self.db = db.create_connection() diff --git a/jarvis/cogs/owner.py b/jarvis/cogs/owner.py index d5b06eb..d4f8249 100644 --- a/jarvis/cogs/owner.py +++ b/jarvis/cogs/owner.py @@ -6,6 +6,12 @@ from jarvis.utils import update class OwnerCog(commands.Cog): + """ + J.A.R.V.I.S. management cog + + Used by admins to control core J.A.R.V.I.S. systems + """ + def __init__(self, bot): self.bot = bot self.admins = get_config().admins diff --git a/jarvis/cogs/util.py b/jarvis/cogs/util.py index ae16d9f..ab0ac8c 100644 --- a/jarvis/cogs/util.py +++ b/jarvis/cogs/util.py @@ -6,6 +6,12 @@ from discord.ext import commands class UtilCog(commands.Cog): + """ + Utility functions for J.A.R.V.I.S. + + Mostly system utility functions, but may change over time + """ + def __init__(self, bot): self.bot = bot self.config = config.get_config() diff --git a/jarvis/utils/__init__.py b/jarvis/utils/__init__.py index f02633b..83d9f74 100644 --- a/jarvis/utils/__init__.py +++ b/jarvis/utils/__init__.py @@ -10,12 +10,21 @@ __all__ = ["field", "db"] def convert_bytesize(bytes: int) -> str: bytes = float(bytes) - sizes = ["B", "KB", "MB", "GB"] + sizes = ["B", "KB", "MB", "GB", "TB", "PB"] size = 0 while bytes >= 1024 and size < len(sizes) - 1: bytes = bytes / 1024 size += 1 - return f"{bytes:0.2f}{sizes[size]}" + return "{:0.3f} {}".format(bytes, sizes[size]) + + +def unconvert_bytesize(size, ending: str): + ending = ending.upper() + sizes = ["B", "KB", "MB", "GB", "TB", "PB"] + if ending == "B": + return size + # Rounding is only because bytes cannot be partial + return round(size * (1024 ** sizes.index(ending))) def get_prefix(bot, message):