Merge pull request #202 from python-discord/jb3/features/face-generator

4 files changed, 90 insertions, 10 deletions

diff --git a/Dockerfile b/Dockerfile
index cbbccc0..40789cf 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,5 +1,8 @@
 FROM --platform=linux/amd64 ghcr.io/owl-corp/python-poetry-base:3.12-slim
 
+# Install build dependencies
+RUN apt-get update && apt-get install -y libmagickwand-dev && apt autoclean && rm -rf /var/lib/apt/lists/*
+
 # Install project dependencies
 WORKDIR /app
 COPY pyproject.toml poetry.lock ./
diff --git a/arthur/exts/systems/system_information.py b/arthur/exts/systems/system_information.py
index fd6384f..03182b1 100644
--- a/arthur/exts/systems/system_information.py
+++ b/arthur/exts/systems/system_information.py
@@ -1,13 +1,16 @@
 """Return system information on our production 9front infrastructure."""
 
 import asyncio
+import io
 import random
 from datetime import UTC, datetime
+from urllib import parse
 
 import aiohttp
-from discord import Member, Message
-from discord.ext.commands import Cog, Context, command
+from discord import File, Member, Message
+from discord.ext.commands import Cog, Context, Converter, command
 from loguru import logger
+from wand.image import Image
 
 from arthur.apis.systems import lib9front
 from arthur.bot import KingArthur
@@ -27,6 +30,20 @@ CORPORATE_FRIENDLY_SMILEYS = (
 )
 
 
+class URLConverter(Converter):
+    """Validate a passed argument is a URL, for use in optional converters that are looking for URLs."""
+
+    async def convert(self, _ctx: Context, argument: str) -> str | None:
+        """Attempt to convert a string to a URL, return the argument if it is a URL, else return None."""
+        try:
+            parsed = parse.urlparse(argument)
+
+            if parsed.scheme in {"http", "https"}:
+                return argument
+        except ValueError:
+            return None
+
+
 class SystemInformation(Cog):
     """Utilities for fetching system information from our 9front infrastructure."""
 
@@ -121,6 +138,51 @@ I enjoy talking to you. Your mind appeals to me. It resembles my own mind except
         program = lib9front.generate_buzzwords(bullshit)
         await ctx.reply(program)
 
+    @command(name="face")
+    async def face(
+        self, ctx: Context, resolution: int | None = 60, *, image_url: URLConverter | None = None
+    ) -> None:
+        """
+        Generate a system-compatible face for the given file.
+
+        If specified, resolution is the integer width and height to use for the generated image, defaulting to 60 (for a 60x60 image).
+
+        The image can be passed in as a URL or attached to the command invocation message.
+        """
+        image_bytes = io.BytesIO()
+
+        if not image_url:
+            if len(ctx.message.attachments) == 0:
+                await ctx.reply(":x: Must upload an image or specify image URL")
+                return
+
+            await ctx.message.attachments[0].save(image_bytes)
+        else:
+            async with aiohttp.ClientSession() as session, session.get(image_url) as resp:
+                if resp.ok:
+                    image_bytes.write(await resp.read())
+                    image_bytes.seek(0)
+                else:
+                    await ctx.reply(
+                        f":x: Could not read remote resource, check it exists. (status `{resp.status}`)"
+                    )
+                    return
+
+        out_bytes = io.BytesIO()
+
+        with Image(file=image_bytes) as img:
+            img.resize(resolution, resolution)
+            img.type = "grayscale"
+            img.kmeans(number_colors=2)
+
+            img.format = "png"
+
+            img.save(file=out_bytes)
+
+        out_bytes.seek(0)
+
+        await ctx.reply(file=File(out_bytes, filename="face.png"))
+
 
 async def setup(bot: KingArthur) -> None:
     """Add cog to bot."""
diff --git a/poetry.lock b/poetry.lock
index fdfa769..e74ccff 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -181,13 +181,13 @@ files = [
 
 [[package]]
 name = "certifi"
-version = "2024.2.2"
+version = "2024.6.2"
 description = "Python package for providing Mozilla's CA Bundle."
 optional = false
 python-versions = ">=3.6"
 files = [
-    {file = "certifi-2024.2.2-py3-none-any.whl", hash = "sha256:dc383c07b76109f368f6106eee2b593b04a011ea4d55f652c6ca24a754d1cdd1"},
-    {file = "certifi-2024.2.2.tar.gz", hash = "sha256:0569859f95fc761b18b45ef421b1290a0f65f147e92a1e5eb3e635f9a5e4e66f"},
+    {file = "certifi-2024.6.2-py3-none-any.whl", hash = "sha256:ddc6c8ce995e6987e7faf5e3f1b02b302836a0e5d98ece18392cb1a36c72ad56"},
+    {file = "certifi-2024.6.2.tar.gz", hash = "sha256:3cd43f1c6fa7dedc5899d69d3ad0398fd018ad1a17fba83ddaf78aa46c747516"},
 ]
 
 [[package]]
@@ -1006,7 +1006,6 @@ files = [
     {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"},
     {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"},
     {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"},
-    {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"},
     {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"},
     {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"},
     {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"},
@@ -1201,13 +1200,13 @@ files = [
 
 [[package]]
 name = "typing-extensions"
-version = "4.12.0"
+version = "4.12.1"
 description = "Backported and Experimental Type Hints for Python 3.8+"
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "typing_extensions-4.12.0-py3-none-any.whl", hash = "sha256:b349c66bea9016ac22978d800cfff206d5f9816951f12a7d0ec5578b0a819594"},
-    {file = "typing_extensions-4.12.0.tar.gz", hash = "sha256:8cbcdc8606ebcb0d95453ad7dc5065e6237b6aa230a31e81d0f440c30fed5fd8"},
+    {file = "typing_extensions-4.12.1-py3-none-any.whl", hash = "sha256:6024b58b69089e5a89c347397254e35f1bf02a907728ec7fee9bf0fe837d203a"},
+    {file = "typing_extensions-4.12.1.tar.gz", hash = "sha256:915f5e35ff76f56588223f15fdd5938f9a1cf9195c0de25130c627e4d597f6d1"},
 ]
 
 [[package]]
@@ -1248,6 +1247,21 @@ docs = ["furo (>=2023.7.26)", "proselint (>=0.13)", "sphinx (>=7.1.2,!=7.3)", "s
 test = ["covdefaults (>=2.3)", "coverage (>=7.2.7)", "coverage-enable-subprocess (>=1)", "flaky (>=3.7)", "packaging (>=23.1)", "pytest (>=7.4)", "pytest-env (>=0.8.2)", "pytest-freezer (>=0.4.8)", "pytest-mock (>=3.11.1)", "pytest-randomly (>=3.12)", "pytest-timeout (>=2.1)", "setuptools (>=68)", "time-machine (>=2.10)"]
 
 [[package]]
+name = "wand"
+version = "0.6.13"
+description = "Ctypes-based simple MagickWand API binding for Python"
+optional = false
+python-versions = "*"
+files = [
+    {file = "Wand-0.6.13-py2.py3-none-any.whl", hash = "sha256:e5dda0ac2204a40c29ef5c4cb310770c95d3d05c37b1379e69c94ea79d7d19c0"},
+    {file = "Wand-0.6.13.tar.gz", hash = "sha256:f5013484eaf7a20eb22d1821aaefe60b50cc329722372b5f8565d46d4aaafcca"},
+]
+
+[package.extras]
+doc = ["Sphinx (>=5.3.0)"]
+test = ["pytest (>=7.2.0)"]
+
+[[package]]
 name = "wcwidth"
 version = "0.2.13"
 description = "Measures the displayed width of unicode strings in a terminal"
@@ -1392,4 +1406,4 @@ multidict = ">=4.0"
 [metadata]
 lock-version = "2.0"
 python-versions = "3.12.*"
-content-hash = "8a180341bc0ba2187b4ac96121f7525dfdef346fafc552301f14a3c00b36bc67"
+content-hash = "7d1314f19dc4d1f4296fe99cd71a234d8f63d2f0b58e8bb4bb98bc0f1b7f7ec4"
diff --git a/pyproject.toml b/pyproject.toml
index 2d716d6..768f932 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -23,6 +23,7 @@ tabulate = { extras = ["widechars"], version = "0.9.0" }
 jishaku = "2.5.2"
 sentry-sdk = "2.3.1"
 humanize = "4.9.0"
+wand = "0.6.13"
 
 [tool.poetry.dev-dependencies]
 pre-commit = "3.7.1"