qr-code-generator/utils/qr_generator.py

"""QR code generation utilities with logo overlay support."""

import io
from typing import Optional, Tuple
from PIL import Image, ImageDraw
import qrcode
import segno


class QRCodeGenerator:
    """Generate QR codes with optional logo overlay and multiple export formats."""

    def __init__(self, url: str, error_correction: str = "H"):
        """
        Initialize QR code generator.

        Args:
            url: The URL to encode in the QR code
            error_correction: Error correction level (L, M, Q, H)
                             H (30%) recommended for logo overlay
        """
        self.url = url
        self.error_correction = self._get_error_correction(error_correction)

    def _get_error_correction(self, level: str) -> int:
        """Convert error correction level string to qrcode constant."""
        levels = {
            "L": qrcode.constants.ERROR_CORRECT_L,
            "M": qrcode.constants.ERROR_CORRECT_M,
            "Q": qrcode.constants.ERROR_CORRECT_Q,
            "H": qrcode.constants.ERROR_CORRECT_H,
        }
        return levels.get(level, qrcode.constants.ERROR_CORRECT_H)

    def generate_qr_image(
        self,
        box_size: int = 10,
        border: int = 4,
        fill_color: str = "black",
        back_color: str = "white"
    ) -> Image.Image:
        """
        Generate QR code as PIL Image.

        Args:
            box_size: Size of each box in pixels
            border: Border size in boxes
            fill_color: QR code color
            back_color: Background color

        Returns:
            PIL Image object of the QR code
        """
        qr = qrcode.QRCode(
            version=1,
            error_correction=self.error_correction,
            box_size=box_size,
            border=border,
        )
        qr.add_data(self.url)
        qr.make(fit=True)

        img = qr.make_image(fill_color=fill_color, back_color=back_color)
        return img.convert("RGB")

    def add_logo(
        self,
        qr_image: Image.Image,
        logo_image: Image.Image,
        logo_size_ratio: float = 0.3
    ) -> Image.Image:
        """
        Add logo to center of QR code.

        Args:
            qr_image: QR code image
            logo_image: Logo image to overlay
            logo_size_ratio: Logo size as ratio of QR code size (default 0.3)

        Returns:
            QR code image with logo overlay
        """
        # Calculate logo size
        qr_width, qr_height = qr_image.size
        logo_max_size = int(min(qr_width, qr_height) * logo_size_ratio)

        # Resize logo while maintaining aspect ratio
        logo_image = logo_image.convert("RGBA")
        logo_image.thumbnail((logo_max_size, logo_max_size), Image.Resampling.LANCZOS)

        # Create white background for logo
        logo_bg_size = int(logo_max_size * 1.1)
        logo_bg = Image.new("RGB", (logo_bg_size, logo_bg_size), "white")

        # Calculate positions
        logo_bg_pos = (
            (qr_width - logo_bg_size) // 2,
            (qr_height - logo_bg_size) // 2
        )
        logo_pos = (
            (logo_bg_size - logo_image.size[0]) // 2,
            (logo_bg_size - logo_image.size[1]) // 2
        )

        # Paste logo background onto QR code
        qr_image.paste(logo_bg, logo_bg_pos)

        # Paste logo with transparency
        absolute_logo_pos = (
            logo_bg_pos[0] + logo_pos[0],
            logo_bg_pos[1] + logo_pos[1]
        )
        qr_image.paste(logo_image, absolute_logo_pos, logo_image)

        return qr_image

    def generate_svg(self, scale: int = 10) -> bytes:
        """
        Generate QR code as SVG.

        Args:
            scale: Scale factor for SVG

        Returns:
            SVG data as bytes
        """
        qr = segno.make(self.url, error="h")
        buffer = io.BytesIO()
        qr.save(buffer, kind="svg", scale=scale, border=4)
        return buffer.getvalue()

    def resize_image(
        self,
        image: Image.Image,
        target_size: Tuple[int, int]
    ) -> Image.Image:
        """
        Resize image to target size.

        Args:
            image: PIL Image to resize
            target_size: Target (width, height) in pixels

        Returns:
            Resized PIL Image
        """
        return image.resize(target_size, Image.Resampling.LANCZOS)

    def export_image(
        self,
        image: Image.Image,
        format: str = "PNG",
        quality: int = 95
    ) -> bytes:
        """
        Export image to specified format.

        Args:
            image: PIL Image to export
            format: Output format (PNG, JPEG, etc.)
            quality: Quality for JPEG (1-100)

        Returns:
            Image data as bytes
        """
        buffer = io.BytesIO()

        if format.upper() == "JPEG":
            # JPEG doesn't support transparency, convert to RGB
            image = image.convert("RGB")
            image.save(buffer, format=format, quality=quality, optimize=True)
        else:
            image.save(buffer, format=format, optimize=True)

        return buffer.getvalue()