esp8266/scripts: Make neopixel/apa102 handle 4bpp LEDs with common code.

The NeoPixel class now handles 4 bytes-per-pixel LEDs (extra byte is
intensity) and arbitrary byte ordering.  APA102 class is now derived from
NeoPixel to reduce code size and support fill() operation.

2 files changed, 22 insertions, 32 deletions

diff --git a/esp8266/scripts/apa102.py b/esp8266/scripts/apa102.py
index 126448cc20..41b7c0485c 100644
--- a/esp8266/scripts/apa102.py
+++ b/esp8266/scripts/apa102.py
@@ -2,27 +2,16 @@
 # MIT license; Copyright (c) 2016 Robert Foss, Daniel Busch
 
 from esp import apa102_write
+from neopixel import NeoPixel
 
-class APA102:
-    def __init__(self, clock_pin, data_pin, n):
-        self.clock_pin = clock_pin
-        self.data_pin = data_pin
-        self.n = n
-        self.buf = bytearray(n * 4)
-
-        self.clock_pin.init(clock_pin.OUT)
-        self.data_pin.init(data_pin.OUT)
 
-    def __setitem__(self, index, val):
-        r, g, b, brightness = val
-        self.buf[index * 4] = r
-        self.buf[index * 4 + 1] = g
-        self.buf[index * 4 + 2] = b
-        self.buf[index * 4 + 3] = brightness
+class APA102(NeoPixel):
+    ORDER = (0, 1, 2, 3)
 
-    def __getitem__(self, index):
-        i = index * 4
-        return self.buf[i], self.buf[i + 1], self.buf[i + 2], self.buf[i + 3]
+    def __init__(self, clock_pin, data_pin, n, bpp=4):
+        super().__init__(data_pin, n, bpp)
+        self.clock_pin = clock_pin
+        self.clock_pin.init(clock_pin.OUT)
 
     def write(self):
-        apa102_write(self.clock_pin, self.data_pin, self.buf)
+        apa102_write(self.clock_pin, self.pin, self.buf)
diff --git a/esp8266/scripts/neopixel.py b/esp8266/scripts/neopixel.py
index 8aa0348680..b13424d7d8 100644
--- a/esp8266/scripts/neopixel.py
+++ b/esp8266/scripts/neopixel.py
@@ -3,29 +3,30 @@
 
 from esp import neopixel_write
 
+
 class NeoPixel:
-    def __init__(self, pin, n):
+    ORDER = (1, 0, 2, 3)
+
+    def __init__(self, pin, n, bpp=3):
         self.pin = pin
         self.n = n
-        self.buf = bytearray(n * 3)
+        self.bpp = bpp
+        self.buf = bytearray(n * bpp)
         self.pin.init(pin.OUT)
 
     def __setitem__(self, index, val):
-        r, g, b = val
-        self.buf[index * 3] = g
-        self.buf[index * 3 + 1] = r
-        self.buf[index * 3 + 2] = b
+        offset = index * self.bpp
+        for i in range(self.bpp):
+            self.buf[offset + self.ORDER[i]] = val[i]
 
     def __getitem__(self, index):
-        i = index * 3
-        return self.buf[i + 1], self.buf[i], self.buf[i + 2]
+        offset = index * self.bpp
+        return tuple(self.buf[offset + self.ORDER[i]]
+                     for i in range(self.bpp))
 
     def fill(self, color):
-        r, g, b = color
-        for i in range(len(self.buf) / 3):
-            self.buf[i * 3] = g
-            self.buf[i * 3 + 1] = r
-            self.buf[i * 3 + 2] = b
+        for i in range(self.n):
+            self[i] = color
 
     def write(self):
         neopixel_write(self.pin, self.buf, True)