pybeamshow/util/color.py

from dataclasses import dataclass
import random
from typing import Generator, Literal, Sequence, Tuple
import pygame as pg


ColorGenerator = Generator[pg.Color, None, None]


def copy_color(source: pg.Color) -> pg.Color:
    return pg.Color(source.r, source.g, source.b, source.a)


@dataclass(frozen=True, slots=True)
class Colors:
    Black = pg.Color(0, 0, 0)
    White = pg.Color(255, 255, 255)
    Red = pg.Color(255, 0, 0)
    Green = pg.Color(0, 255, 0)
    Blue = pg.Color(0, 0, 255)
    Cyan = pg.Color(0, 255, 255)
    Yellow = pg.Color(255, 255, 0)
    Magenta = pg.Color(255, 0, 255)


def color_static(color: pg.Color) -> ColorGenerator:
    while True:
        yield color


def color_wheel(hue: int = 0, increase: int = 1, repeat: int = 1) -> ColorGenerator:
    color = copy_color(Colors.Red)
    h, s, l, a = color.hsla
    h = hue % 360

    while True:
        color.hsla = h, s, l, a
        for _ in range(repeat):
            yield color
        h = (h + increase) % 360


def color_randomize() -> ColorGenerator:
    color = copy_color(Colors.Red)
    h, s, l, a = color.hsla
    color.hsla = random.randrange(0, 360 // 5) * 5, s, l, a

    while True:
        yield color
        color.hsla = random.randrange(0, 360 // 5) * 5, s, l, a


def color_fadeout(initial_color: pg.Color, fade_speed: float) -> ColorGenerator:
    color = copy_color(initial_color)
    h, s, l, a = color.hsla
    l_float = float(l)

    while True:
        yield color
        color.hsla = h, s, int(l_float), a
        l_float -= fade_speed
        if l_float < 0:
            l_float = 0


def color_fade(
    initial_color: pg.Color, end_color: pg.Color, duration: int
) -> ColorGenerator:
    inc = 1.0 / duration
    factor = 0.0
    col1 = copy_color(initial_color)
    col2 = copy_color(end_color)
    while True:
        yield col1.lerp(col2, factor)
        factor += inc
        if factor > 1:
            factor = 1


def color_cycle(seq: Sequence[pg.Color]) -> ColorGenerator:
    it = iter(seq)
    while True:
        yield next(it)


def color_shuffle(seq: Sequence[pg.Color]) -> ColorGenerator:
    while True:
        yield random.choice(seq)


def color_darken(color: pg.Color, factor: float) -> pg.Color:
    h, s, l, a = color.hsla
    new_color = pg.Color(0, 0, 0, 255)
    new_color.hsla = h, s, l * factor, a
    return new_color


def color_strobe(
    color: ColorGenerator,
    rate: Tuple[int, int],
    phase: int = 0,
    flash_color: ColorGenerator = color_static(Colors.Black),
) -> ColorGenerator:
    counter = phase
    while True:
        if counter < rate[0]:
            yield next(flash_color)
        else:
            yield next(color)
        counter = (counter + 1) % rate[1]


def rainbow_surface(
    image: pg.Surface,
    orientation: Literal["h", "v"] = "h",
    hue: int = 0,
    increase: int = 1,
):
    wheel = color_wheel(hue, increase)
    if orientation == "h":
        h = image.get_height()
        for x in range(image.get_width() // 5):
            pg.draw.line(image, next(wheel), (x * 5, 0), (x * 5, h), width=5)
    elif orientation == "v":
        w = image.get_width()
        for y in range(image.get_height() // 5):
            pg.draw.line(image, next(wheel), (0, y * 5), (w, y * 5), width=5)