Add PixelSubStrip class to simplify working with multiple segments

library/rpi_ws281x/rpi_ws281x.py

@@ -105,6 +105,15 @@ def __init__(self, num, pin, freq_hz=800000, dma=10, invert=False,
         # Grab the led data array.
         self._led_data = _LED_Data(self._channel, num)
 
+        # Create a PixelSubStrip and delegate these methods to it
+        self.main_strip = self.PixelSubStrip(self, 0, num=num)
+        self.setPixelColor = self.main_strip.setPixelColor
+        self.setPixelColorRGB = self.main_strip.setPixelColorRGB
+        self.setBrightness = self.main_strip.setBrightness
+        self.getBrightness = self.main_strip.getBrightness
+        self.getPixels = self.main_strip.getPixels
+        self.getPixelColor = self.main_strip.getPixelColor
+
         # Substitute for __del__, traps an exit condition and cleans up properly
         atexit.register(self._cleanup)
 
@@ -137,55 +146,109 @@ def show(self):
             str_resp = ws.ws2811_get_return_t_str(resp)
             raise RuntimeError('ws2811_render failed with code {0} ({1})'.format(resp, str_resp))
 
-    def setPixelColor(self, n, color):
-        """Set LED at position n to the provided 24-bit color value (in RGB order).
-        """
-        self._led_data[n] = color
-
-    def setPixelColorRGB(self, n, red, green, blue, white=0):
-        """Set LED at position n to the provided red, green, and blue color.
-        Each color component should be a value from 0 to 255 (where 0 is the
-        lowest intensity and 255 is the highest intensity).
-        """
-        self.setPixelColor(n, Color(red, green, blue, white))
-
-    def getBrightness(self):
-        return ws.ws2811_channel_t_brightness_get(self._channel)
-
-    def setBrightness(self, brightness):
-        """Scale each LED in the buffer by the provided brightness.  A brightness
-        of 0 is the darkest and 255 is the brightest.
-        """
-        ws.ws2811_channel_t_brightness_set(self._channel, brightness)
-
-    def getPixels(self):
-        """Return an object which allows access to the LED display data as if
-        it were a sequence of 24-bit RGB values.
-        """
-        return self._led_data
-
     def numPixels(self):
         """Return the number of pixels in the display."""
         return ws.ws2811_channel_t_count_get(self._channel)
 
-    def getPixelColor(self, n):
-        """Get the 24-bit RGB color value for the LED at position n."""
-        return self._led_data[n]
+    def createPixelSubStrip(self, first, last=None, num=None):
+        """Create a PixelSubStrip starting with pixel `first`
+        Either specify the `num` of pixels or the `last` pixel.
+
+        All the methods of a PixelSubStrip are available on PixelStrip
+        objects.
 
-    def getPixelColorRGB(self, n):
-        c = lambda: None
-        setattr(c, 'r', self._led_data[n] >> 16 & 0xff)
-        setattr(c, 'g', self._led_data[n] >> 8  & 0xff)
-        setattr(c, 'b', self._led_data[n]    & 0xff)
-        return c
+        Note: PixelSubStrips are not prevented from overlappping
+        """
+        if last:
+            if last > self.numPixels():
+                raise self.InvalidStrip(f"Too many pixels ({last})."
+                                        f"Strip only has {self.numPixels()}.")
+            return self.PixelSubStrip(self, first, last=last)
+        if num:
+            if first + num > self.numPixels():
+                raise self.InvalidStrip(f"Too many pixels ({first + num})."
+                                        f"Strip only has {self.numPixels()}.")
+            return self.PixelSubStrip(self, first, num=num)
+        raise self.InvalidStrip("Need num or last to create a PixelSubStrip")
+
+    class InvalidStrip(Exception):
+        pass
 
-    def getPixelColorRGBW(self, n):
-        c = lambda: None
-        setattr(c, 'w', self._led_data[n] >> 24 & 0xff)
-        setattr(c, 'r', self._led_data[n] >> 16 & 0xff)
-        setattr(c, 'g', self._led_data[n] >> 8  & 0xff)
-        setattr(c, 'b', self._led_data[n]    & 0xff)
-        return c
+    class PixelSubStrip:
+        """A PixelSubStrip handles a subset of the pixels in a PixelStrip
+
+        strip = PixelStrip(...)
+        strip1 = strip.createPixelSubStrip(0, num=10)  # controls first 10 pixels
+        strip2 = strip.createPixelSubStrip(10, num=10)  # controls next 10 pixels
+        """
+        def __init__(self, strip, first, last=None, num=None):
+            self.strip = strip
+            self.first = first
+            if last:
+                self.last = last
+                self.num = last - first
+            elif num:
+                self.last = first + num
+                self.num = num
+            else:
+                raise self.InvalidStrip("Must specify number or last pixel to "
+                                        "create a PixelSubStrip")
+
+        def setPixelColor(self, n, color):
+            """Set LED at position n to the provided 24-bit color value (in RGB order).
+            """
+            self.strip._led_data[self.first + n] = color
+
+        def setPixelColorRGB(self, n, red, green, blue, white=0):
+            """Set LED at position n to the provided red, green, and blue color.
+            Each color component should be a value from 0 to 255 (where 0 is the
+            lowest intensity and 255 is the highest intensity).
+            """
+            # No translation to n - do that in the called method
+            self.setPixelColor(n, Color(red, green, blue, white))
+
+        def getBrightness(self):
+            return ws.ws2811_channel_t_brightness_get(self.strip._channel)
+
+        def setBrightness(self, brightness):
+            """Scale each LED in the buffer by the provided brightness.  A brightness
+            of 0 is the darkest and 255 is the brightest.
+
+            This method affects all pixels in all PixelSubStrips.
+            """
+            ws.ws2811_channel_t_brightness_set(self.strip._channel, brightness)
+
+        def getPixels(self):
+            """Return an object which allows access to the LED display data as if
+            it were a sequence of 24-bit RGB values.
+            """
+            return self.strip._led_data
+
+        def getPixelColor(self, n):
+            """Get the 24-bit RGB color value for the LED at position n."""
+            return self.strip._led_data[self.first + n]
+
+        def getPixelColorRGB(self, n):
+            c = lambda: None
+            setattr(c, 'r', self.strip._led_data[self.first + n] >> 16 & 0xff)
+            setattr(c, 'g', self.strip._led_data[self.first + n] >> 8  & 0xff)
+            setattr(c, 'b', self.strip._led_data[self.first + n]    & 0xff)
+            return c
+
+        def numPixels(self):
+            """Return the number of pixels in the strip."""
+            return self.num
+
+        def getPixelColorRGBW(self, n):
+            c = lambda: None
+            setattr(c, 'w', self.strip._led_data[self.first + n] >> 24 & 0xff)
+            setattr(c, 'r', self.strip._led_data[self.first + n] >> 16 & 0xff)
+            setattr(c, 'g', self.strip._led_data[self.first + n] >> 8  & 0xff)
+            setattr(c, 'b', self.strip._led_data[self.first + n]    & 0xff)
+            return c
+
+        def show(self):
+            self.strip.show()
 
 # Shim for back-compatibility
 class Adafruit_NeoPixel(PixelStrip):