feat: add iteration over Palette and FlavorColors (#30)

Author: backwardspy

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),

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]

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