Merge pull request #931

[BUG] ColorsGenerator never returns random_colors

Author: Leguark

gempy/core/color_generator.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+from functools import cached_property
 from typing import Optional
 
 import numpy as np
@@ -36,11 +37,10 @@ def __init__(self):
     @staticmethod
     def _random_hexcolor() -> str:
         """Generates a random hex color string."""
-        return "#"+str(hex(np.random.randint(0, 16777215))).lstrip("0x")
+        return f"#{np.random.randint(0, 0xffffff):06x}"
 
     def regenerate_color_palette(self, seaborn_palettes: Optional[list[str]] = None):
         try:
-            
             import seaborn as sns
             seaborn_installed = True
         except ImportError:
@@ -54,7 +54,7 @@ def regenerate_color_palette(self, seaborn_palettes: Optional[list[str]] = None)
         elif seaborn_palettes and not seaborn_installed:
             raise ImportError("Seaborn is not installed. Please install it to use color palettes.")
         else:
-            hex_colors = self._gempy_default_colors
+            hex_colors = list(self._gempy_default_colors)
 
         self.hex_colors = hex_colors
 
@@ -65,14 +65,18 @@ def __iter__(self) -> 'ColorsGenerator':
 
     def __next__(self) -> str:
         """Generator that yields the next color."""
-        for color in self.hex_colors:
-            result = self.hex_colors[self._index]
-            self._index += 1
-            return result
+        color = self.up_next()
+        self._index += 1
+        del self._next_color
+        return color
 
-        while True:
-            return self._random_hexcolor()
-    
     def up_next(self) -> str:
         """Returns the next color without incrementing the index."""
-        return self.hex_colors[self._index]
+        return self._next_color
+
+    @cached_property
+    def _next_color(self) -> str:
+        if self._index < len(self.hex_colors):
+            return self.hex_colors[self._index]
+
+        return self._random_hexcolor()

test/test_core/test_color_generator.py

@@ -1,10 +1,10 @@
-from gempy.core.color_generator import ColorsGenerator
+import gempy.core.color_generator as cg
 
 
 def test_color_generator():
 
     # create an instance of ColorsGenerator
-    color_iterator: iter = ColorsGenerator()
+    color_iterator: iter = cg.ColorsGenerator()
 
     # generate color dictionary (with optional seaborn palettes if you want)
 
@@ -13,4 +13,32 @@ def test_color_generator():
     # now you can get colors one by one
     print(next(color_iterator))  # prints first color
     print(next(color_iterator))  # prints second color
+
+    left_colors = color_iterator._gempy_default_colors[2:]
+
+    assert left_colors == [
+        color for color, _ in zip(color_iterator, left_colors)
+    ]
+
+    # get a random color
+    assert color_iterator.up_next() == next(color_iterator)
+
+    # override the random color generator
+    cg.np.random.randint = lambda a, b: 0x00ff42
+
+    assert color_iterator._random_hexcolor() == "#00ff42"
+    assert next(color_iterator) == "#00ff42"
+
+    # trigger up_next method for the cache
+    assert color_iterator.up_next() == "#00ff42"
+
+    cg.np.random.randint = lambda a, b: 0x000000
+
+    # ensure that the up_next method still returns the same color
+    assert color_iterator.up_next() == "#00ff42"
+
+    # trigger the cache
+    assert next(color_iterator) == "#00ff42"
+    assert next(color_iterator) == "#000000"
+
     # etc...