diff --git a/README.md b/README.md index 8d82e83..a836cf6 100644 --- a/README.md +++ b/README.md @@ -34,6 +34,31 @@ PALETTE.mocha.colors.teal.rgb The `Palette` data structure matches [the palette JSON](https://github.com/catppuccin/palette/blob/main/palette.json). +### Iteration + +Both `Palette` and `FlavorColors` can be iterated to yield flavors and colors respectively: + +```python +for flavor in PALETTE: + print(flavor.name) + +# Latte +# Frappé +# Macchiato +# Mocha + +for color in PALETTE.latte.colors: + print(f"{color.name}: {color.hex}") + +# Rosewater: #f2d5cf +# Flamingo: #eebebe +# Pink: #f4b8e4 +# ... +# Base: #303446 +# Mantle: #292c3c +# Crust: #232634 +``` + ### dataclasses `Palette`, `Flavor`, `Color` et cetera are all [`dataclasses`](https://docs.python.org/3/library/dataclasses.html), diff --git a/catppuccin/models.py b/catppuccin/models.py index 3ea2385..686d62d 100644 --- a/catppuccin/models.py +++ b/catppuccin/models.py @@ -1,5 +1,6 @@ """Dataclass definitions for the Catppuccin palette data structure.""" from dataclasses import dataclass +from typing import Iterator @dataclass(frozen=True) @@ -63,6 +64,37 @@ class FlavorColors: mantle: Color crust: Color + def __iter__(self) -> Iterator[Color]: + """Iterate over colors in the flavor.""" + yield from [ + self.rosewater, + self.flamingo, + self.pink, + self.mauve, + self.red, + self.maroon, + self.peach, + self.yellow, + self.green, + self.teal, + self.sky, + self.sapphire, + self.blue, + self.lavender, + self.text, + self.subtext1, + self.subtext0, + self.overlay2, + self.overlay1, + self.overlay0, + self.surface2, + self.surface1, + self.surface0, + self.base, + self.mantle, + self.crust, + ] + @dataclass(frozen=True) class Flavor: @@ -91,3 +123,7 @@ class Palette: frappe: Flavor macchiato: Flavor mocha: Flavor + + def __iter__(self) -> Iterator[Flavor]: + """Iterate over flavors in the palette.""" + yield from [self.latte, self.frappe, self.macchiato, self.mocha] diff --git a/tests/test_catppuccin.py b/tests/test_catppuccin.py index 9f82821..0559053 100644 --- a/tests/test_catppuccin.py +++ b/tests/test_catppuccin.py @@ -6,3 +6,45 @@ def test_some_colors() -> None: assert PALETTE.latte.colors.mauve.rgb.r == 136 assert PALETTE.latte.colors.mauve.rgb.g == 57 assert PALETTE.latte.colors.mauve.rgb.b == 239 + + +def test_iterate_palette() -> None: + order = [PALETTE.latte, PALETTE.frappe, PALETTE.macchiato, PALETTE.mocha] + for i, flavor in enumerate(PALETTE): + assert order[i] == flavor + assert list(PALETTE) == order + + +def test_iterate_flavor_colors() -> None: + colors = PALETTE.latte.colors + order = [ + colors.rosewater, + colors.flamingo, + colors.pink, + colors.mauve, + colors.red, + colors.maroon, + colors.peach, + colors.yellow, + colors.green, + colors.teal, + colors.sky, + colors.sapphire, + colors.blue, + colors.lavender, + colors.text, + colors.subtext1, + colors.subtext0, + colors.overlay2, + colors.overlay1, + colors.overlay0, + colors.surface2, + colors.surface1, + colors.surface0, + colors.base, + colors.mantle, + colors.crust, + ] + for i, color in enumerate(colors): + assert order[i] == color + assert list(colors) == order