Add PixelSubStrip class to simplify working with multiple segments

.gitignore

@@ -15,3 +15,5 @@ library/debian/
 dist/
 __pycache__
 .DS_Store
+.coverage
+coverage.xml

library/pytest.ini

@@ -0,0 +1,2 @@
+[pytest]
+addopts = --cov=rpi_ws281x --cov-report xml

library/rpi_ws281x/__init__.py

@@ -1,5 +1,5 @@
 # New canonical package, to support `import rpi_ws281x`
-from .rpi_ws281x import PixelStrip, Adafruit_NeoPixel, Color, RGBW, ws
+from .rpi_ws281x import PixelStrip, InvalidStrip, Adafruit_NeoPixel, Color, RGBW, ws
 from _rpi_ws281x import *
 
 __version__ = '5.0.0'

library/rpi_ws281x/rpi_ws281x.py

@@ -11,6 +11,10 @@ def __new__(self, r, g=None, b=None, w=None):
         else:
             if w is None:
                 w = 0
+            if g is None:
+                g = 0
+            if b is None:
+                b = 0
             return int.__new__(self, (w << 24) | (r << 16) | (g << 8) | b)
 
     @property
@@ -48,6 +52,9 @@ def __init__(self, num, pin, freq_hz=800000, dma=10, invert=False,
         800khz), dma, the DMA channel to use (default 10), invert, a boolean
         specifying if the signal line should be inverted (default False), and
         channel, the PWM channel to use (defaults to 0).
+
+        All the methods of a PixelSubStrip are available on PixelStrip
+        objects.
         """
 
         if gamma is None:
@@ -89,6 +96,18 @@ def __init__(self, num, pin, freq_hz=800000, dma=10, invert=False,
 
         self.size = 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
+        self.getPixelColorRGB = self.main_strip.getPixelColorRGB
+        self.getPixelColorRGBW = self.main_strip.getPixelColorRGBW
+        self.off = self.main_strip.off
+
         # Substitute for __del__, traps an exit condition and cleans up properly
         atexit.register(self._cleanup)
 
@@ -106,20 +125,30 @@ def __getitem__(self, pos):
 
     def __setitem__(self, pos, value):
         """Set the 24-bit RGB color value at the provided position or slice of
-        positions.
+        positions. If value is a slice it is zip()'ed with pos to set as many
+        leds as there are values.
         """
         # Handle if a slice of positions are passed in by setting the appropriate
         # LED data values to the provided value.
+        # Cast to int() as value may be a numpy non-int value.
         if isinstance(pos, slice):
-            for n in range(*pos.indices(self.size)):
-                ws.ws2811_led_set(self._channel, n, value)
+            try:
+                for n, c in zip(range(*pos.indices(self.size)), value):
+                    ws.ws2811_led_set(self._channel, n, int(c))
+            except TypeError:
+                for n in range(*pos.indices(self.size)):
+                    ws.ws2811_led_set(self._channel, n, int(value))
         # Else assume the passed in value is a number to the position.
         else:
-            return ws.ws2811_led_set(self._channel, pos, value)
+            return ws.ws2811_led_set(self._channel, pos, int(value))
 
     def __len__(self):
         return ws.ws2811_channel_t_count_get(self._channel)
 
+    def numPixels(self):
+        """Return the number of pixels in the display."""
+        return len(self)
+
     def _cleanup(self):
         # Clean up memory used by the library when not needed anymore.
         if self._leds is not None:
@@ -149,46 +178,148 @@ 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[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 createPixelSubStrip(self, first, last=None, num=None):
+        """Create a PixelSubStrip starting with pixel `first`
+        Either specify the `num` of pixels or the `last` pixel.
 
-    def getBrightness(self):
-        return ws.ws2811_channel_t_brightness_get(self._channel)
+        All the methods of a PixelSubStrip are available on PixelStrip
+        objects.
 
-    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.
+        Note: PixelSubStrips are not prevented from overlappping
         """
-        ws.ws2811_channel_t_brightness_set(self._channel, brightness)
+        return self.PixelSubStrip(self, first, last=last, num=num)
 
-    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[:]
 
-    def numPixels(self):
-        """Return the number of pixels in the display."""
-        return len(self)
+    class PixelSubStrip:
+        """A PixelSubStrip handles a subset of the pixels in a PixelStrip
 
-    def getPixelColor(self, n):
-        """Get the 24-bit RGB color value for the LED at position n."""
-        return self[n]
+        strip = PixelStrip(...)
+        strip1 = strip.createPixelSubStrip(0, num=10)  # controls first 10 pixels
+        strip2 = strip.createPixelSubStrip(10, num=10)  # controls next 10 pixels
 
-    def getPixelColorRGB(self, n):
-        return RGBW(self[n])
+        strip2[5] will access the 15th pixel
+        """
 
-    def getPixelColorRGBW(self, n):
-        return RGBW(self[n])
+        def __init__(self, strip, first, last=None, num=None):
+            self.strip = strip
+            if first < 0:
+                raise InvalidStrip(f"First pixel is negative ({first}).")
+            if first > len(strip):
+                raise InvalidStrip(f"First pixel is too big ({first})."
+                                   f"Strip only has {len(strip)}.")
+            self.first = first
+            if last:
+                if last < 0:
+                    raise InvalidStrip(f"Last pixel is negative ({last}).")
+                if last > len(strip):
+                    raise InvalidStrip(f"Too many pixels ({last})."
+                                       f"Strip only has {len(strip)}.")
+                self.last = last
+                self.num = last - first
+            elif num:
+                if num < 0:
+                    raise InvalidStrip(f"number of pixels is negative ({num}).")
+                if first + num > len(strip):
+                    raise InvalidStrip(f"Too many pixels (last would be {first + num})."
+                                       f"Strip only has {len(strip)}.")
+                self.last = first + num
+                self.num = num
+            else:
+                raise InvalidStrip("Must specify number or last pixel to "
+                                   "create a PixelSubStrip")
+
+        def __len__(self):
+            return self.num
+
+        def _adjust_pos(self, pos):
+            # create an adjusted pos, either a slice or index
+            if isinstance(pos, slice):
+                if pos.start and pos.start < 0:
+                    apos_start = self.first + self.num + pos.start
+                else:
+                    apos_start = (0 if pos.start is None else pos.start) + self.first
+
+                if pos.stop and pos.stop < 0:
+                    apos_stop = pos.stop + self.last
+                else:
+                    apos_stop = (self.num if pos.stop is None else pos.stop) + self.first
+                apos = slice(apos_start,
+                             apos_stop,
+                             pos.step)
+                return apos
+            if pos < 0:
+                return self.num + pos + self.first
+            return pos + self.first
+
+
+        def __getitem__(self, pos):
+            """Return the 24-bit RGB color value at the provided position or slice
+            of positions.
+            """
+            return self.strip[self._adjust_pos(pos)]
+
+        def __setitem__(self, pos, value):
+            """Set the 24-bit RGB color value at the provided position or slice of
+            positions. If value is a slice it is zip()'ed with pos to set as many
+            leds as there are values.
+            """
+            self.strip[self._adjust_pos(pos)] = value
+
+        def setPixelColor(self, n, color):
+            """Set LED at position n to the provided 24-bit color value (in RGB order).
+            If n is a slice then color can be a value which is repeated for all leds
+            or a slice of values which are applied to the leds.
+            """
+            self[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).
+            """
+            # Translation to n done in setPixelColor
+            self[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[:]
+
+        def numPixels(self):
+            """Return the number of pixels in the strip."""
+            return self.num
+
+        def getPixelColor(self, n):
+            """Get the 24-bit RGB color value for the LED at position n."""
+            return self[n]
+
+        def getPixelColorRGB(self, n):
+            return RGBW(self[n])
+
+        def getPixelColorRGBW(self, n):
+            return RGBW(self[n])
+
+        def show(self):
+            self.strip.show()
+
+        def off(self):
+            self[:] = 0
+            self.strip.show()
+
+class InvalidStrip(Exception):
+    pass
 
 # Shim for back-compatibility
 class Adafruit_NeoPixel(PixelStrip):

library/tests/conftest.py

@@ -39,6 +39,7 @@ def ws2811_led_get(ch, n):
 @pytest.fixture(scope='function', autouse=False)
 def _rpi_ws281x():
     _mock_rpi_ws281x.ws2811_init.return_value = 0
+    _mock_rpi_ws281x.ws2811_render.return_value = 0
     sys.modules['_rpi_ws281x'] = _mock_rpi_ws281x
 
     yield _mock_rpi_ws281x