diff --git a/src/emotigrad/__init__.py b/src/emotigrad/__init__.py index 7aeb712..b83639b 100644 --- a/src/emotigrad/__init__.py +++ b/src/emotigrad/__init__.py @@ -1,3 +1,4 @@ from .base import EmotionalOptimizer +from .colors import ColoredPrinter, colorize, create_colored_print_fn -__all__ = ["EmotionalOptimizer"] +__all__ = ["EmotionalOptimizer", "ColoredPrinter", "colorize", "create_colored_print_fn"] diff --git a/src/emotigrad/colors.py b/src/emotigrad/colors.py new file mode 100644 index 0000000..bf037a7 --- /dev/null +++ b/src/emotigrad/colors.py @@ -0,0 +1,219 @@ +# src/emotigrad/colors.py +"""Colored console output support for emotigrad. + +This module provides ANSI color codes and utility functions for +colorful console output. It includes personality-specific color schemes. +""" + +from __future__ import annotations + +from typing import Optional + +# ANSI escape codes for colors +RESET = "\033[0m" +BOLD = "\033[1m" +DIM = "\033[2m" +ITALIC = "\033[3m" +UNDERLINE = "\033[4m" + +# Foreground colors +BLACK = "\033[30m" +RED = "\033[31m" +GREEN = "\033[32m" +YELLOW = "\033[33m" +BLUE = "\033[34m" +MAGENTA = "\033[35m" +CYAN = "\033[36m" +WHITE = "\033[37m" + +# Bright foreground colors +BRIGHT_BLACK = "\033[90m" +BRIGHT_RED = "\033[91m" +BRIGHT_GREEN = "\033[92m" +BRIGHT_YELLOW = "\033[93m" +BRIGHT_BLUE = "\033[94m" +BRIGHT_MAGENTA = "\033[95m" +BRIGHT_CYAN = "\033[96m" +BRIGHT_WHITE = "\033[97m" + +# Background colors +BG_BLACK = "\033[40m" +BG_RED = "\033[41m" +BG_GREEN = "\033[42m" +BG_YELLOW = "\033[43m" +BG_BLUE = "\033[44m" +BG_MAGENTA = "\033[45m" +BG_CYAN = "\033[46m" +BG_WHITE = "\033[47m" + + +def colorize(text: str, *codes: str) -> str: + """Apply ANSI color codes to text. + + Parameters + ---------- + text : str + The text to colorize. + *codes : str + ANSI color codes to apply. + + Returns + ------- + str + The colorized text with reset code at the end. + + Examples + -------- + >>> colorize("Hello", GREEN, BOLD) + '\\033[32m\\033[1mHello\\033[0m' + """ + if not codes: + return text + return "".join(codes) + text + RESET + + +def strip_colors(text: str) -> str: + """Remove ANSI color codes from text. + + Parameters + ---------- + text : str + Text potentially containing ANSI codes. + + Returns + ------- + str + Text with all ANSI codes removed. + """ + import re + + ansi_escape = re.compile(r"\033\[[0-9;]*m") + return ansi_escape.sub("", text) + + +# Personality-specific color schemes +PERSONALITY_COLORS = { + "wholesome": (GREEN, BOLD), + "sassy": (MAGENTA, BOLD), + "quiet": (DIM,), + "nervous": (YELLOW,), + "chaotic": (BRIGHT_MAGENTA, BOLD), + "arrogant": (CYAN, ITALIC), + "tired": (DIM, ITALIC), + "hype": (BRIGHT_YELLOW, BOLD), + "academic": (BLUE,), + "pirate": (BRIGHT_RED, BOLD), + "zen": (BRIGHT_CYAN, DIM), +} + + +def get_personality_colors(personality_name: str) -> tuple: + """Get the color codes for a specific personality. + + Parameters + ---------- + personality_name : str + Name of the personality. + + Returns + ------- + tuple + Tuple of ANSI color codes for the personality. + Returns empty tuple if personality not found. + """ + return PERSONALITY_COLORS.get(personality_name.lower(), ()) + + +class ColoredPrinter: + """A printer that applies personality-specific colors to output. + + This can be used as the `print_fn` parameter in EmotionalOptimizer + to get colored output based on the personality being used. + + Parameters + ---------- + personality_name : str, optional + Name of the personality to use for coloring. + If None, no colors will be applied. + enabled : bool, optional + Whether to enable colored output. Default is True. + Set to False to disable colors (useful for non-TTY output). + + Examples + -------- + >>> from emotigrad import EmotionalOptimizer + >>> from emotigrad.colors import ColoredPrinter + >>> + >>> printer = ColoredPrinter("hype") + >>> emo_opt = EmotionalOptimizer(optimizer, personality="hype", print_fn=printer) + """ + + def __init__( + self, + personality_name: Optional[str] = None, + enabled: bool = True, + ) -> None: + self.personality_name = personality_name + self.enabled = enabled + self._colors = ( + get_personality_colors(personality_name) if personality_name else () + ) + + def __call__(self, message: str) -> None: + """Print the message with personality-specific colors. + + Parameters + ---------- + message : str + The message to print. + """ + if self.enabled and self._colors: + print(colorize(message, *self._colors)) + else: + print(message) + + def set_personality(self, personality_name: str) -> None: + """Update the personality and associated colors. + + Parameters + ---------- + personality_name : str + Name of the new personality. + """ + self.personality_name = personality_name + self._colors = get_personality_colors(personality_name) + + +def create_colored_print_fn( + personality_name: Optional[str] = None, + enabled: bool = True, +) -> callable: + """Create a colored print function for a specific personality. + + This is a convenience function that returns a callable + suitable for use as EmotionalOptimizer's print_fn. + + Parameters + ---------- + personality_name : str, optional + Name of the personality for coloring. + enabled : bool, optional + Whether to enable colors. + + Returns + ------- + callable + A function that prints with colors. + + Examples + -------- + >>> from emotigrad import EmotionalOptimizer + >>> from emotigrad.colors import create_colored_print_fn + >>> + >>> emo_opt = EmotionalOptimizer( + ... optimizer, + ... personality="academic", + ... print_fn=create_colored_print_fn("academic"), + ... ) + """ + return ColoredPrinter(personality_name, enabled) diff --git a/tests/test_colors.py b/tests/test_colors.py new file mode 100644 index 0000000..cad67aa --- /dev/null +++ b/tests/test_colors.py @@ -0,0 +1,198 @@ +# tests/test_colors.py +"""Tests for the colored output module.""" + +import pytest + +from emotigrad.colors import ( + BLUE, + BOLD, + ColoredPrinter, + GREEN, + PERSONALITY_COLORS, + RED, + RESET, + colorize, + create_colored_print_fn, + get_personality_colors, + strip_colors, +) + + +class TestColorize: + """Tests for the colorize function.""" + + def test_colorize_applies_single_code(self): + """colorize should apply a single color code.""" + result = colorize("Hello", GREEN) + assert result.startswith(GREEN) + assert result.endswith(RESET) + assert "Hello" in result + + def test_colorize_applies_multiple_codes(self): + """colorize should apply multiple color codes.""" + result = colorize("Hello", GREEN, BOLD) + assert GREEN in result + assert BOLD in result + assert result.endswith(RESET) + + def test_colorize_no_codes_returns_original(self): + """colorize with no codes should return original text.""" + result = colorize("Hello") + assert result == "Hello" + + def test_colorize_empty_string(self): + """colorize should handle empty strings.""" + result = colorize("", GREEN) + assert result == GREEN + RESET + + +class TestStripColors: + """Tests for the strip_colors function.""" + + def test_strip_colors_removes_codes(self): + """strip_colors should remove all ANSI codes.""" + colored = colorize("Hello World", GREEN, BOLD) + result = strip_colors(colored) + assert result == "Hello World" + + def test_strip_colors_plain_text(self): + """strip_colors should handle plain text.""" + result = strip_colors("Hello World") + assert result == "Hello World" + + def test_strip_colors_multiple_colors(self): + """strip_colors should remove multiple color sequences.""" + text = f"{RED}Red{RESET} and {BLUE}Blue{RESET}" + result = strip_colors(text) + assert result == "Red and Blue" + + +class TestGetPersonalityColors: + """Tests for get_personality_colors function.""" + + def test_get_colors_for_known_personality(self): + """Should return colors for known personalities.""" + colors = get_personality_colors("wholesome") + assert len(colors) > 0 + assert colors == PERSONALITY_COLORS["wholesome"] + + def test_get_colors_case_insensitive(self): + """Should be case insensitive.""" + assert get_personality_colors("HYPE") == get_personality_colors("hype") + assert get_personality_colors("Academic") == get_personality_colors("academic") + + def test_get_colors_unknown_personality(self): + """Should return empty tuple for unknown personality.""" + colors = get_personality_colors("nonexistent") + assert colors == () + + def test_all_registered_personalities_have_colors(self): + """All personalities in PERSONALITY_COLORS should have color tuples.""" + for name, colors in PERSONALITY_COLORS.items(): + assert isinstance(colors, tuple) + assert len(colors) > 0 + + +class TestColoredPrinter: + """Tests for the ColoredPrinter class.""" + + def test_colored_printer_with_personality(self, capsys): + """ColoredPrinter should apply colors when personality is set.""" + printer = ColoredPrinter("wholesome") + printer("Test message") + + captured = capsys.readouterr() + # Should contain the color codes + assert GREEN in captured.out or "Test message" in captured.out + + def test_colored_printer_without_personality(self, capsys): + """ColoredPrinter without personality should print plain text.""" + printer = ColoredPrinter() + printer("Test message") + + captured = capsys.readouterr() + assert "Test message" in captured.out + + def test_colored_printer_disabled(self, capsys): + """ColoredPrinter should not apply colors when disabled.""" + printer = ColoredPrinter("wholesome", enabled=False) + printer("Test message") + + captured = capsys.readouterr() + # Should be plain text + assert strip_colors(captured.out).strip() == "Test message" + + def test_colored_printer_set_personality(self): + """set_personality should update the colors.""" + printer = ColoredPrinter("wholesome") + original_colors = printer._colors + + printer.set_personality("hype") + assert printer._colors != original_colors + assert printer._colors == PERSONALITY_COLORS["hype"] + + +class TestCreateColoredPrintFn: + """Tests for create_colored_print_fn function.""" + + def test_create_returns_callable(self): + """Should return a callable.""" + fn = create_colored_print_fn("wholesome") + assert callable(fn) + + def test_created_fn_prints_with_colors(self, capsys): + """Created function should print with colors.""" + fn = create_colored_print_fn("hype") + fn("Test message") + + captured = capsys.readouterr() + assert "Test message" in captured.out + + def test_created_fn_respects_enabled_flag(self, capsys): + """Created function should respect enabled flag.""" + fn = create_colored_print_fn("hype", enabled=False) + fn("Test message") + + captured = capsys.readouterr() + assert strip_colors(captured.out).strip() == "Test message" + + +class TestIntegrationWithEmotionalOptimizer: + """Integration tests with EmotionalOptimizer.""" + + def test_colored_printer_as_print_fn(self): + """ColoredPrinter should work as EmotionalOptimizer's print_fn.""" + import torch + + from emotigrad import EmotionalOptimizer + + model = torch.nn.Linear(2, 1) + optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + + messages = [] + + class CapturingColoredPrinter(ColoredPrinter): + def __call__(self, message): + messages.append(message) + super().__call__(message) + + printer = CapturingColoredPrinter("wholesome") + + emo_opt = EmotionalOptimizer( + optimizer, + personality="wholesome", + print_fn=printer, + enabled=True, + ) + + x = torch.randn(2, 2) + y = torch.randn(2, 1) + preds = model(x) + loss = (preds - y).pow(2).mean() + + emo_opt.zero_grad() + loss.backward() + emo_opt.step(loss=loss.item()) + + # Should have captured at least one message + assert len(messages) >= 1