feat: cover generation with Pollinations.ai and text fallback
Made-with: Cursor
This commit is contained in:
123
src/cover.py
Normal file
123
src/cover.py
Normal file
@@ -0,0 +1,123 @@
|
||||
import logging
|
||||
import os
|
||||
from datetime import date
|
||||
from io import BytesIO
|
||||
from urllib.parse import quote
|
||||
|
||||
import requests
|
||||
from PIL import Image as PILImage, ImageDraw, ImageFont
|
||||
|
||||
import config
|
||||
from src.images import _resize_to_fit
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _get_font(size: int) -> ImageFont.FreeTypeFont:
|
||||
try:
|
||||
return ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", size)
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
return ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", size)
|
||||
except OSError:
|
||||
pass
|
||||
return ImageFont.load_default()
|
||||
|
||||
|
||||
def generate_text_cover(
|
||||
output_dir: str,
|
||||
week_start: date,
|
||||
week_end: date,
|
||||
headlines: list[str],
|
||||
) -> str:
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
img = PILImage.new("RGB", (800, 480), color="white")
|
||||
draw = ImageDraw.Draw(img)
|
||||
|
||||
title_font = _get_font(36)
|
||||
date_font = _get_font(20)
|
||||
headline_font = _get_font(16)
|
||||
|
||||
draw.text((400, 30), "Plymouth Independent", fill="black",
|
||||
font=title_font, anchor="mt")
|
||||
|
||||
date_str = f"Week of {week_start.strftime('%b %d')} \u2013 {week_end.strftime('%b %d, %Y')}"
|
||||
draw.text((400, 80), date_str, fill="gray", font=date_font, anchor="mt")
|
||||
|
||||
draw.line([(50, 110), (750, 110)], fill="black", width=2)
|
||||
|
||||
y = 130
|
||||
for i, headline in enumerate(headlines[:8]):
|
||||
if y > 440:
|
||||
break
|
||||
prefix = f"\u2022 {headline}"
|
||||
if len(prefix) > 70:
|
||||
prefix = prefix[:67] + "..."
|
||||
draw.text((60, y), prefix, fill="black", font=headline_font)
|
||||
y += 35
|
||||
|
||||
filename = f"cover-{week_start.isoformat()}-text.jpg"
|
||||
path = os.path.join(output_dir, filename)
|
||||
img.save(path, format="JPEG", progressive=False, quality=90)
|
||||
return path
|
||||
|
||||
|
||||
def generate_ai_cover(
|
||||
output_dir: str,
|
||||
week_start: date,
|
||||
week_end: date,
|
||||
headlines: list[str],
|
||||
) -> str:
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
top_headlines = ", ".join(headlines[:3])
|
||||
prompt = (
|
||||
f"Newspaper front page illustration for Plymouth Massachusetts local news. "
|
||||
f"Headlines: {top_headlines}. "
|
||||
f"Classic broadsheet newspaper style, black and white ink drawing, editorial illustration."
|
||||
)
|
||||
url = config.POLLINATIONS_URL.format(prompt=quote(prompt))
|
||||
|
||||
try:
|
||||
response = requests.get(url, timeout=60)
|
||||
response.raise_for_status()
|
||||
|
||||
img = PILImage.open(BytesIO(response.content))
|
||||
if img.mode != "RGB":
|
||||
img = img.convert("RGB")
|
||||
img = _resize_to_fit(img)
|
||||
|
||||
draw = ImageDraw.Draw(img)
|
||||
title_font = _get_font(28)
|
||||
date_font = _get_font(16)
|
||||
|
||||
draw.text((img.width // 2, 15), "Plymouth Independent",
|
||||
fill="white", font=title_font, anchor="mt",
|
||||
stroke_width=2, stroke_fill="black")
|
||||
|
||||
date_str = f"Week of {week_start.strftime('%b %d')} \u2013 {week_end.strftime('%b %d, %Y')}"
|
||||
draw.text((img.width // 2, 50), date_str,
|
||||
fill="white", font=date_font, anchor="mt",
|
||||
stroke_width=1, stroke_fill="black")
|
||||
|
||||
filename = f"cover-{week_start.isoformat()}-ai.jpg"
|
||||
path = os.path.join(output_dir, filename)
|
||||
img.save(path, format="JPEG", progressive=False, quality=90)
|
||||
return path
|
||||
|
||||
except Exception as e:
|
||||
logger.error("AI cover generation failed, falling back to text: %s", e)
|
||||
return generate_text_cover(output_dir, week_start, week_end, headlines)
|
||||
|
||||
|
||||
def generate_cover(
|
||||
method: str,
|
||||
output_dir: str,
|
||||
week_start: date,
|
||||
week_end: date,
|
||||
headlines: list[str],
|
||||
) -> str:
|
||||
if method == "ai":
|
||||
return generate_ai_cover(output_dir, week_start, week_end, headlines)
|
||||
return generate_text_cover(output_dir, week_start, week_end, headlines)
|
||||
77
tests/test_cover.py
Normal file
77
tests/test_cover.py
Normal file
@@ -0,0 +1,77 @@
|
||||
import os
|
||||
from datetime import date
|
||||
from unittest.mock import patch, MagicMock
|
||||
from io import BytesIO
|
||||
from PIL import Image as PILImage
|
||||
|
||||
from src.cover import generate_text_cover, generate_ai_cover, generate_cover
|
||||
|
||||
|
||||
def test_text_cover_creates_jpeg(tmp_path):
|
||||
path = generate_text_cover(
|
||||
output_dir=str(tmp_path),
|
||||
week_start=date(2026, 4, 6),
|
||||
week_end=date(2026, 4, 12),
|
||||
headlines=["Article One", "Article Two", "Article Three"],
|
||||
)
|
||||
|
||||
assert os.path.exists(path)
|
||||
img = PILImage.open(path)
|
||||
assert img.format == "JPEG"
|
||||
assert img.width <= 800
|
||||
assert img.height <= 480
|
||||
|
||||
|
||||
def test_ai_cover_creates_jpeg(tmp_path):
|
||||
fake_img = PILImage.new("RGB", (800, 480), color="blue")
|
||||
buf = BytesIO()
|
||||
fake_img.save(buf, format="JPEG")
|
||||
fake_bytes = buf.getvalue()
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.content = fake_bytes
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
|
||||
with patch("src.cover.requests.get", return_value=mock_response):
|
||||
path = generate_ai_cover(
|
||||
output_dir=str(tmp_path),
|
||||
week_start=date(2026, 4, 6),
|
||||
week_end=date(2026, 4, 12),
|
||||
headlines=["Test Headline"],
|
||||
)
|
||||
|
||||
assert os.path.exists(path)
|
||||
img = PILImage.open(path)
|
||||
assert img.format == "JPEG"
|
||||
assert img.width <= 800
|
||||
assert img.height <= 480
|
||||
|
||||
|
||||
def test_ai_cover_falls_back_on_failure(tmp_path):
|
||||
with patch("src.cover.requests.get", side_effect=Exception("API down")):
|
||||
path = generate_ai_cover(
|
||||
output_dir=str(tmp_path),
|
||||
week_start=date(2026, 4, 6),
|
||||
week_end=date(2026, 4, 12),
|
||||
headlines=["Test"],
|
||||
)
|
||||
|
||||
assert os.path.exists(path)
|
||||
img = PILImage.open(path)
|
||||
assert img.format == "JPEG"
|
||||
|
||||
|
||||
def test_generate_cover_dispatches(tmp_path):
|
||||
with patch("src.cover.generate_ai_cover") as mock_ai:
|
||||
mock_ai.return_value = "/fake/ai.jpg"
|
||||
result = generate_cover("ai", str(tmp_path), date(2026, 4, 6),
|
||||
date(2026, 4, 12), ["A"])
|
||||
assert result == "/fake/ai.jpg"
|
||||
mock_ai.assert_called_once()
|
||||
|
||||
with patch("src.cover.generate_text_cover") as mock_text:
|
||||
mock_text.return_value = "/fake/text.jpg"
|
||||
result = generate_cover("text", str(tmp_path), date(2026, 4, 6),
|
||||
date(2026, 4, 12), ["A"])
|
||||
assert result == "/fake/text.jpg"
|
||||
mock_text.assert_called_once()
|
||||
Reference in New Issue
Block a user