diff --git a/polybot/image.py b/polybot/image.py index 32179e0..3b8b57b 100644 --- a/polybot/image.py +++ b/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) diff --git a/polybot/service.py b/polybot/service.py index a522a7a..0648dde 100644 --- a/polybot/service.py +++ b/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() diff --git a/tests/images/sample.png b/tests/images/sample.png new file mode 100644 index 0000000..41112da Binary files /dev/null and b/tests/images/sample.png differ diff --git a/tests/test_image.py b/tests/test_image.py new file mode 100644 index 0000000..298a056 --- /dev/null +++ b/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