Fix comet length and simplification

adafruit_led_animation/animation/comet.py

@@ -45,7 +45,7 @@
 """
 
 from adafruit_led_animation.animation import Animation
-from adafruit_led_animation.color import BLACK
+from adafruit_led_animation.color import BLACK, calculate_intensity
 
 
 class Comet(Animation):
@@ -62,7 +62,7 @@ class Comet(Animation):
     :param bool bounce: Comet will bounce back and forth. Defaults to ``True``.
     """
 
-    # pylint: disable=too-many-arguments
+    # pylint: disable=too-many-arguments,too-many-instance-attributes
     def __init__(
         self,
         pixel_object,
@@ -75,75 +75,59 @@ def __init__(
     ):
         if tail_length == 0:
             tail_length = len(pixel_object) // 4
-        else:
-            tail_length = max(2, min(tail_length, len(pixel_object)))
-        self._tail_length = tail_length
-        self._color_step = 0.9 / tail_length
-        self._color_offset = 0.1
-        self._comet_colors = None
-        self._reverse_comet_colors = None
-        self._initial_reverse = reverse
         self.reverse = reverse
         self.bounce = bounce
+        self._initial_reverse = reverse
+        self._tail_length = tail_length
+        self._color_step = 0.95 / tail_length
+        self._comet_colors = None
         self._computed_color = color
-        self._generator = self._comet_generator()
+        self._num_pixels = len(pixel_object)
+        self._direction = -1 if reverse else 1
+        self._left_side = -self._tail_length
+        self._right_side = self._num_pixels
+        self._tail_start = 0
+        self.reset()
         super().__init__(pixel_object, speed, color, name=name)
 
     on_cycle_complete_supported = True
 
     def _recompute_color(self, color):
-        pass
+        self._comet_recompute_color(color)
 
     def _comet_recompute_color(self, color):
-        self._comet_colors = [BLACK] + [
-            [
-                int(color[rgb] * ((n * self._color_step) + self._color_offset))
-                for rgb in range(len(color))
-            ]
-            for n in range(self._tail_length - 1)
-        ]
-        self._reverse_comet_colors = list(reversed(self._comet_colors))
+        self._comet_colors = [BLACK]
+        for n in range(self._tail_length):
+            self._comet_colors.append(
+                calculate_intensity(color, n * self._color_step + 0.05)
+            )
         self._computed_color = color
 
-    def _get_range(self, num_pixels):
+    def draw(self):
+        colors = self._comet_colors
         if self.reverse:
-            return range(num_pixels, -self._tail_length - 1, -1)
-        return range(-self._tail_length, num_pixels + 1)
-
-    def _comet_generator(self):
-        num_pixels = len(self.pixel_object)
-        cycle_passes = 0
-        while True:
-            if self._color != self._computed_color or not self._comet_colors:
-                self._comet_recompute_color(self._color)
-            colors = self._reverse_comet_colors if self.reverse else self._comet_colors
-            for start in self._get_range(num_pixels):
-
-                if start + self._tail_length < num_pixels:
-                    end = self._tail_length
-                else:
-                    end = num_pixels - start
-                if start <= 0:
-                    num_visible = self._tail_length + start
-                    self.pixel_object[0:num_visible] = colors[
-                        self._tail_length - num_visible :
-                    ]
-                else:
-                    self.pixel_object[start : start + end] = colors[0:end]
-                yield
-            cycle_passes += 1
+            colors = reversed(colors)
+        for pixel_no, color in enumerate(colors):
+            draw_at = self._tail_start + pixel_no
+            if draw_at < 0 or draw_at >= self._num_pixels:
+                continue
+            self.pixel_object[draw_at] = color
+
+        self._tail_start += self._direction
+
+        if self._tail_start < self._left_side or self._tail_start >= self._right_side:
             if self.bounce:
                 self.reverse = not self.reverse
-            if not self.bounce or cycle_passes == 2:
+                self._direction = -self._direction
+            if self.reverse == self._initial_reverse:
                 self.cycle_complete = True
-                cycle_passes = 0
-
-    def draw(self):
-        next(self._generator)
 
     def reset(self):
         """
         Resets to the first color.
         """
-        self._generator = self._comet_generator()
         self.reverse = self._initial_reverse
+        if self.reverse:
+            self._tail_start = self._num_pixels + self._tail_length + 1
+        else:
+            self._tail_start = -self._tail_length - 1

adafruit_led_animation/animation/rainbowcomet.py

@@ -45,7 +45,7 @@
 """
 
 from adafruit_led_animation.animation.comet import Comet
-from adafruit_led_animation.color import colorwheel, BLACK
+from adafruit_led_animation.color import colorwheel, BLACK, calculate_intensity
 
 
 class RainbowComet(Comet):
@@ -60,6 +60,7 @@ class RainbowComet(Comet):
     :param bool reverse: Animates the comet in the reverse order. Defaults to ``False``.
     :param bool bounce: Comet will bounce back and forth. Defaults to ``True``.
     :param int colorwheel_offset: Offset from start of colorwheel (0-255).
+    :param int step: Colorwheel step (defaults to automatic).
     """
 
     # pylint: disable=too-many-arguments
@@ -71,30 +72,27 @@ def __init__(
         reverse=False,
         bounce=False,
         colorwheel_offset=0,
+        step=0,
         name=None,
     ):
-        self._colorwheel_is_tuple = isinstance(colorwheel(0), tuple)
+        if step == 0:
+            self._colorwheel_step = int(256 / tail_length)
+        else:
+            self._colorwheel_step = step
         self._colorwheel_offset = colorwheel_offset
-
         super().__init__(pixel_object, speed, 0, tail_length, reverse, bounce, name)
 
-    def _calc_brightness(self, n, color):
-        brightness = (n * self._color_step) + self._color_offset
-        if not self._colorwheel_is_tuple:
-            color = (color & 0xFF, ((color & 0xFF00) >> 8), (color >> 16))
-        return [int(i * brightness) for i in color]
-
     def _comet_recompute_color(self, color):
-        factor = int(256 / self._tail_length)
-        self._comet_colors = [BLACK] + [
-            self._calc_brightness(
-                n,
-                colorwheel(
-                    int((n * factor) + self._color_offset + self._colorwheel_offset)
-                    % 256
-                ),
+        self._comet_colors = [BLACK]
+        for n in range(self._tail_length):
+            invert = self._tail_length - n - 1
+            self._comet_colors.append(
+                calculate_intensity(
+                    colorwheel(
+                        int((invert * self._colorwheel_step) + self._colorwheel_offset)
+                        % 256
+                    ),
+                    n * self._color_step + 0.05,
+                )
             )
-            for n in range(self._tail_length - 1)
-        ]
-        self._reverse_comet_colors = list(reversed(self._comet_colors))
         self._computed_color = color

adafruit_led_animation/color.py

@@ -64,3 +64,40 @@ def colorwheel(pos):
             return 0, int(255 - pos * 3), int(pos * 3)
         pos -= 170
         return int(pos * 3), 0, int(255 - (pos * 3))
+
+
+def calculate_intensity(color, intensity=1.0):
+    """
+    Takes a RGB[W] color tuple and adjusts the intensity.
+    :param float intensity:
+    :param color: color value (tuple, list or int)
+    :return: color
+    """
+    # Note: This code intentionally avoids list comprehensions and intermediate variables
+    # for an approximately 2x performance gain.
+    if isinstance(color, int):
+        return (
+            (int((color & 0xFF0000) * intensity) & 0xFF0000)
+            | (int((color & 0xFF00) * intensity) & 0xFF00)
+            | (int((color & 0xFF) * intensity) & 0xFF)
+        )
+
+    if len(color) == 3:
+        return (
+            int(color[0] * intensity),
+            int(color[1] * intensity),
+            int(color[2] * intensity),
+        )
+    if len(color) == 4 and isinstance(color[3], float):
+        return (
+            int(color[0] * intensity),
+            int(color[1] * intensity),
+            int(color[2] * intensity),
+            color[3],
+        )
+    return (
+        int(color[0] * intensity),
+        int(color[1] * intensity),
+        int(color[2] * intensity),
+        int(color[3] * intensity),
+    )