Rebase onto upstream (a4d95fd)
#12
@@ -10,6 +10,9 @@ from pathlib import Path
|
||||
from loguru import logger
|
||||
from PIL import Image
|
||||
|
||||
# Telegram API photo size limit (6MB)
|
||||
TELEGRAM_PHOTO_SIZE_LIMIT = 6_000_000
|
||||
|
||||
try:
|
||||
import magic
|
||||
HAS_MAGIC = True
|
||||
@@ -102,7 +105,7 @@ def is_heic_format(path: str) -> bool:
|
||||
return ext in (".heic", ".heif")
|
||||
|
||||
|
||||
async def optimize_image(path: str, max_bytes: int = 6_000_000) -> bytes:
|
||||
def optimize_image(path: str, max_bytes: int = TELEGRAM_PHOTO_SIZE_LIMIT) -> bytes:
|
||||
"""
|
||||
Optimize image to fit under size limit.
|
||||
|
||||
@@ -134,32 +137,27 @@ async def optimize_image(path: str, max_bytes: int = 6_000_000) -> bytes:
|
||||
# Convert to RGB (HEIC → JPEG)
|
||||
if img.mode != "RGB":
|
||||
img = img.convert("RGB")
|
||||
return await _optimize_jpeg(img, max_bytes)
|
||||
return _optimize_jpeg(img, max_bytes)
|
||||
|
||||
# PNG with alpha channel - preserve it
|
||||
if img.mode == "RGBA" or img.mode == "LA":
|
||||
return await _optimize_png(img, max_bytes)
|
||||
return _optimize_png(img, max_bytes)
|
||||
|
||||
# Everything else → convert to JPEG and optimize
|
||||
if img.mode != "RGB":
|
||||
img = img.convert("RGB")
|
||||
return await _optimize_jpeg(img, max_bytes)
|
||||
return _optimize_jpeg(img, max_bytes)
|
||||
|
||||
|
||||
async def _optimize_jpeg(img: Image.Image, max_bytes: int) -> bytes:
|
||||
def _optimize_jpeg(img: Image.Image, max_bytes: int) -> bytes:
|
||||
"""Optimize JPEG with size/quality grid."""
|
||||
sizes = [2048, 1536, 1280, 1024, 800]
|
||||
qualities = [80, 70, 60, 50, 40]
|
||||
|
||||
original_width, original_height = img.size
|
||||
|
||||
for size in sizes:
|
||||
# Skip if image is already smaller
|
||||
if max(original_width, original_height) <= size:
|
||||
resized = img
|
||||
else:
|
||||
# Resize maintaining aspect ratio
|
||||
resized = img.copy()
|
||||
# Always copy to avoid mutation issues
|
||||
resized = img.copy()
|
||||
if max(img.size) > size:
|
||||
resized.thumbnail((size, size), Image.Resampling.LANCZOS)
|
||||
|
||||
for quality in qualities:
|
||||
@@ -174,19 +172,15 @@ async def _optimize_jpeg(img: Image.Image, max_bytes: int) -> bytes:
|
||||
raise ValueError(f"Cannot optimize image under {max_bytes} bytes")
|
||||
|
||||
|
||||
async def _optimize_png(img: Image.Image, max_bytes: int) -> bytes:
|
||||
def _optimize_png(img: Image.Image, max_bytes: int) -> bytes:
|
||||
"""Optimize PNG while preserving alpha channel."""
|
||||
compress_levels = [6, 7, 8, 9]
|
||||
sizes = [2048, 1536, 1280, 1024, 800]
|
||||
|
||||
original_width, original_height = img.size
|
||||
|
||||
for size in sizes:
|
||||
# Skip if image is already smaller
|
||||
if max(original_width, original_height) <= size:
|
||||
resized = img
|
||||
else:
|
||||
resized = img.copy()
|
||||
# Always copy to avoid mutation issues
|
||||
resized = img.copy()
|
||||
if max(img.size) > size:
|
||||
resized.thumbnail((size, size), Image.Resampling.LANCZOS)
|
||||
|
||||
for compress_level in compress_levels:
|
||||
@@ -198,10 +192,13 @@ async def _optimize_png(img: Image.Image, max_bytes: int) -> bytes:
|
||||
return data
|
||||
|
||||
# Fallback: try converting to JPEG if still too large
|
||||
if img.mode == "RGBA":
|
||||
if img.mode in ("RGBA", "LA"):
|
||||
# Create white background
|
||||
background = Image.new("RGB", img.size, (255, 255, 255))
|
||||
background.paste(img, mask=img.split()[3]) # Use alpha as mask
|
||||
return await _optimize_jpeg(background, max_bytes)
|
||||
if img.mode == "RGBA":
|
||||
background.paste(img, mask=img.split()[3]) # Use alpha as mask
|
||||
else: # LA (grayscale + alpha)
|
||||
background.paste(img.convert("L"), mask=img.split()[1])
|
||||
return _optimize_jpeg(background, max_bytes)
|
||||
|
||||
raise ValueError(f"Cannot optimize PNG under {max_bytes} bytes")
|
||||
|
||||
@@ -1,9 +1,7 @@
|
||||
"""Tests for Telegram media handling."""
|
||||
|
||||
import io
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
|
||||
@@ -116,8 +114,7 @@ def test_is_heic_format():
|
||||
assert is_heic_format("photo.jpg") is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_optimize_image_jpeg_quality():
|
||||
def test_optimize_image_jpeg_quality(tmp_path):
|
||||
"""Test JPEG optimization reduces size with quality ladder."""
|
||||
from nanobot.channels.telegram_media import optimize_image
|
||||
|
||||
@@ -129,27 +126,21 @@ async def test_optimize_image_jpeg_quality():
|
||||
original_size = len(original_bytes)
|
||||
|
||||
# Write to temp file
|
||||
import tempfile
|
||||
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
|
||||
f.write(original_bytes)
|
||||
temp_path = f.name
|
||||
temp_file = tmp_path / "test.jpg"
|
||||
temp_file.write_bytes(original_bytes)
|
||||
|
||||
try:
|
||||
# Optimize to 1MB max
|
||||
optimized = await optimize_image(temp_path, max_bytes=1_000_000)
|
||||
# Optimize to 1MB max
|
||||
optimized = optimize_image(str(temp_file), max_bytes=1_000_000)
|
||||
|
||||
# Should be smaller than original
|
||||
assert len(optimized) < original_size
|
||||
# Should be under limit
|
||||
assert len(optimized) <= 1_000_000
|
||||
# Should still be valid JPEG
|
||||
assert optimized.startswith(b'\xff\xd8\xff')
|
||||
finally:
|
||||
Path(temp_path).unlink()
|
||||
# Should be smaller than original
|
||||
assert len(optimized) < original_size
|
||||
# Should be under limit
|
||||
assert len(optimized) <= 1_000_000
|
||||
# Should still be valid JPEG
|
||||
assert optimized.startswith(b'\xff\xd8\xff')
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_optimize_image_png_preserve_alpha():
|
||||
def test_optimize_image_png_preserve_alpha(tmp_path):
|
||||
"""Test PNG with alpha channel is preserved."""
|
||||
from nanobot.channels.telegram_media import optimize_image
|
||||
|
||||
@@ -159,19 +150,15 @@ async def test_optimize_image_png_preserve_alpha():
|
||||
img.save(buf, format="PNG")
|
||||
original_bytes = buf.getvalue()
|
||||
|
||||
import tempfile
|
||||
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
|
||||
f.write(original_bytes)
|
||||
temp_path = f.name
|
||||
# Write to temp file
|
||||
temp_file = tmp_path / "test.png"
|
||||
temp_file.write_bytes(original_bytes)
|
||||
|
||||
try:
|
||||
optimized = await optimize_image(temp_path, max_bytes=5_000_000)
|
||||
optimized = optimize_image(str(temp_file), max_bytes=5_000_000)
|
||||
|
||||
# Should still be PNG (PNG magic bytes)
|
||||
assert optimized.startswith(b'\x89PNG')
|
||||
# Should still be PNG (PNG magic bytes)
|
||||
assert optimized.startswith(b'\x89PNG')
|
||||
|
||||
# Load and verify alpha channel preserved
|
||||
img_opt = Image.open(io.BytesIO(optimized))
|
||||
assert img_opt.mode == "RGBA"
|
||||
finally:
|
||||
Path(temp_path).unlink()
|
||||
# Load and verify alpha channel preserved
|
||||
img_opt = Image.open(io.BytesIO(optimized))
|
||||
assert img_opt.mode == "RGBA"
|
||||
|
||||
Reference in New Issue
Block a user