Make RainbowChase work, and make Chase code easier to work with via generators

adafruit_led_animation/animation.py

@@ -72,6 +72,7 @@ class Animation:
         self.color = color  # Triggers _recompute_color
         self.done_cycle_handler = None
         self.name = name
+        self.counter = 0
 
     def __str__(self):
         return "<Animation %s: %s>" % (self.__class__.__name__, self.name)
@@ -91,6 +92,8 @@ class Animation:
             return False
 
         self.draw()
+        self.counter += 1
+
         if self.peers:
             for peer in self.peers:
                 peer.draw()
@@ -555,10 +558,10 @@ class Chase(Animation):
         self._overflow = len(pixel_object) % self._repeat_width
         self._direction = 1 if not reverse else -1
         self._reverse = reverse
-        self._n = 0
+        self._offset = 0
 
         def _resetter():
-            self._n = 0
+            self._offset = 0
             self._reverse = reverse
             self._direction = 1 if not reverse else -1
 
@@ -579,27 +582,47 @@ class Chase(Animation):
         self._direction = -1 if self._reverse else 1
 
     def draw(self):
-        self.pixel_object.fill((0, 0, 0))
-        for i in range(self._size):
-            n = (self._n + i) % self._repeat_width
-            num = len(self.pixel_object[n :: self._repeat_width])
-            self.pixel_object[n :: self._repeat_width] = [
-                self.group_color(n) for n in range(num)
-            ]
-        _n = (self._n + self._direction) % self._repeat_width
-        if _n < self._n:
-            self._cycle_done()
-        self._n = _n
-        self.show()
 
-    def group_color(self, n):  # pylint: disable=unused-argument
+        def bar_colors():
+            bar_no = 0
+            for i in range(self._offset, 0, -1):
+                if i > self._spacing:
+                    yield self.bar_color(bar_no, i)
+                else:
+                    yield self.space_color(bar_no, i)
+                    bar_no = 1
+            while True:
+                for bar_pixel in range(self._size):
+                    yield self.bar_color(bar_no, bar_pixel)
+                for space_pixel in range(self._spacing):
+                    yield self.space_color(bar_no, space_pixel)
+                bar_no += 1
+
+        colorgen = bar_colors()
+        self.pixel_object[:] = [next(colorgen) for _ in self.pixel_object]
+
+        if self._offset == 0:
+            self._cycle_done()
+        self._offset = (self._offset + self._direction) % self._repeat_width
+
+    def bar_color(self, n, pixel_no=0):  # pylint: disable=unused-argument
         """
-        Generate the color for the n'th group
+        Generate the color for the n'th bar_color in the Chase
 
         :param n: The pixel group to get the color for
+        :param pixel_no: Which pixel in the group to get the color for
         """
         return self.color
 
+    def space_color(self, n, pixel_no=0):
+        """
+        Generate the spacing color for the n'th bar_color in the Chase
+
+        :param n: The pixel group to get the spacing color for
+        :param pixel_no: Which pixel in the group to get the spacing color for
+        """
+        return 0
+
     def reset(self):
         """
         Reset the animation.
@@ -607,6 +630,37 @@ class Chase(Animation):
         self._reset()
 
 
+class RainbowChase(Chase):
+    """
+    Chase pixels in one direction, like a theater marquee but with rainbows!
+
+    :param pixel_object: The initialised LED object.
+    :param float speed: Animation speed rate in seconds, e.g. ``0.1``.
+    :param color: Animation color in ``(r, g, b)`` tuple, or ``0x000000`` hex format.
+    :param size: Number of pixels to turn on in a row.
+    :param spacing: Number of pixels to turn off in a row.
+    :param reverse: Reverse direction of movement.
+    :param wheel_step: How many colors to skip in `colorwheel` per bar (default 8)
+    """
+    # pylint: disable=too-many-arguments
+    def __init__(self, pixel_object, speed, color, size=2, spacing=3, reverse=False, name=None,
+                 wheel_step=8):
+        self._num_colors = 256 // wheel_step
+        self._colors = [colorwheel(n % 256) for n in range(0, 512, wheel_step)]
+        self._color_idx = 0
+        super(RainbowChase, self).__init__(pixel_object, speed, color, size, spacing, reverse, name)
+
+    def bar_color(self, n, *_):
+        return self._colors[self._color_idx - n]
+
+    def _cycle_done(self):
+        self._color_idx = (self._color_idx + self._direction) % len(self._colors)
+        super(RainbowChase, self)._cycle_done()
+
+    def show(self):
+        super(RainbowChase, self).show()
+
+
 class AnimationSequence:
     """
     A sequence of Animations to run in sequence, looping forever.