Handle image pixel limits, add tests for image resizing

polybot/image.py

@@ -42,23 +42,32 @@ def __init__(
         self.mime_type = mime_type
         self.description = description
 
-    def resize_to_target(self, target_bytes: int) -> "Image":
-        """Resize the image to a target size in bytes. Returns a new Image object.
+    def resize_to_target(
+        self, target_bytes: int, target_pixels: Optional[int] = None
+    ) -> "Image":
+        """Resize the image to a target maximum size in bytes and (optionally) pixels. Returns a new Image object.
         This is required for Bluesky's silly image size limit.
         """
+
         original_bytes = len(self.data)
-        if original_bytes < target_bytes:
+        if target_pixels is None and original_bytes < target_bytes:
             return self
 
         img = PILImage.open(BytesIO(self.data))
-
         margin = 0.9
         new_bytes = original_bytes
+        new_pixels = original_pixels = img.size[0] * img.size[1]
+        output_bytes = self.data
+
+        if target_pixels is None:
+            target_pixels = original_pixels
+
+        while new_bytes > target_bytes or new_pixels > target_pixels:
+            new_pixels = int(original_pixels * (target_bytes * margin / original_bytes))
+            if target_pixels is not None:
+                new_pixels = min(new_pixels, target_pixels)
 
-        while new_bytes > target_bytes:
-            current_pixels = img.size[0] * img.size[1]
-            new_pixels = int(current_pixels * (target_bytes * margin / original_bytes))
-            ratio = (new_pixels / current_pixels) ** 0.5
+            ratio = (new_pixels / original_pixels) ** 0.5
             new_size = (int(img.width * ratio), int(img.height * ratio))
 
             new_img = img.resize(new_size)

polybot/service.py

@@ -27,6 +27,7 @@ class Service(object):
     max_length = None  # type: int
     max_length_image = None  # type: int
     max_image_size: int = int(10e6)
+    max_image_pixels: Optional[int] = None
 
     def __init__(self, config, live: bool) -> None:
         self.log = logging.getLogger(__name__)
@@ -59,7 +60,10 @@ def post(
         lon: Optional[float] = None,
         in_reply_to_id=None,
     ):
-        images = [i.resize_to_target(self.max_image_size) for i in images]
+        images = [
+            i.resize_to_target(self.max_image_size, self.max_image_pixels)
+            for i in images
+        ]
         if self.live:
             if wrap:
                 return self.do_wrapped(status, images, lat, lon, in_reply_to_id)
@@ -238,9 +242,10 @@ def auth(self):
         )
         self.log.info("Connected to %s at %s", self.software, base_url)
         self.log.info(
-            "Max post length: %d chars, max image size: %d MB",
+            "Max post length: %d chars, max image size: %d MB, max image pixels: %d",
             self.max_length,
             self.max_image_size / 1024 / 1024,
+            self.max_image_pixels,
         )
 
     def fetch_endpoint(self, path):
@@ -252,7 +257,7 @@ def fetch_endpoint(self, path):
             return None
         return res.json()
 
-    def update_instance_info(self):
+    def get_node_software(self):
         data = self.fetch_endpoint("/.well-known/nodeinfo")
         if not data:
             return None
@@ -270,7 +275,14 @@ def update_instance_info(self):
             return None
 
         data = res.json()
-        self.software = data.get("software", {}).get("name")
+        return data.get("software", {}).get("name")
+
+    def update_instance_info(self):
+        """Fetch and save details about the instance we're connecting to, including software type
+        and post size limits.
+        """
+
+        self.software = self.get_node_software()
 
         instance_info = self.fetch_endpoint("/api/v1/instance")
         if not instance_info:
@@ -283,6 +295,15 @@ def update_instance_info(self):
         except Exception:
             image_size = self.max_image_size
 
+        try:
+            image_pixels = int(
+                instance_info["configuration"]["media_attachments"][
+                    "image_matrix_limit"
+                ]
+            )
+        except Exception:
+            image_pixels = self.max_image_pixels
+
         try:
             max_length = int(
                 instance_info["configuration"]["statuses"]["max_characters"]
@@ -293,6 +314,7 @@ def update_instance_info(self):
         self.max_image_size = image_size
         self.max_length = max_length
         self.max_length_image = max_length
+        self.max_image_pixels = image_pixels
 
     def setup(self):
         print()

tests/test_image.py

@@ -0,0 +1,29 @@
+from polybot.image import Image
+from io import BytesIO
+from PIL import Image as PILImage
+from pathlib import Path
+
+
+def count_pixels(image):
+    img = PILImage.open(BytesIO(image.data))
+    return img.size[0] * img.size[1]
+
+
+def test_image_resize():
+    img = Image(
+        path=Path(__file__).parent / "images" / "sample.png", mime_type="image/png"
+    )
+
+    original_size = len(img.data)
+
+    resized = img.resize_to_target(original_size + 1)
+    assert img == resized
+
+    for target_bytes in (1000000, 1500000):
+        resized = img.resize_to_target(target_bytes)
+        assert img.mime_type == resized.mime_type
+        assert len(resized.data) <= target_bytes
+
+    for target_pixels in (1200 * 1200, 1000 * 1000):
+        resized = img.resize_to_target(5000000, target_pixels)
+        assert count_pixels(resized) <= target_pixels