fix sleep and wake for built-in displays

kmk/extensions/oled.py

@@ -84,12 +84,21 @@ def during_bootup(self, width, height, rotation):
     def deinit(self):
         raise NotImplementedError
 
+    def sleep(self):
+        raise NotImplementedError
+
+    def wake(self):
+        raise NotImplementedError
+
 
 # Intended for displays with drivers built into CircuitPython
 # that can be used directly without manual initialization
 class BuiltInDisplay(OledDisplayType):
-    def __init__(self, display=None):
+    def __init__(self, display=None, sleep_command=None, wake_command=None):
         self.display = display
+        self.sleep_command = sleep_command
+        self.wake_command = wake_command
+        self.is_awake = True
 
     def during_bootup(self, width, height, rotation):
         self.display.rotation = rotation
@@ -98,6 +107,12 @@ def during_bootup(self, width, height, rotation):
     def deinit(self):
         return
 
+    def sleep(self):
+        self.display.bus.send(self.sleep_command, b"")
+
+    def wake(self):
+        self.display.bus.send(self.wake_command, b"")
+
 
 class SSD1306(OledDisplayType):
     def __init__(self, i2c=None, sda=None, scl=None, device_address=0x3C):
@@ -111,17 +126,25 @@ def __init__(self, i2c=None, sda=None, scl=None, device_address=0x3C):
     def during_bootup(self, width, height, rotation):
         import adafruit_displayio_ssd1306
 
-        return adafruit_displayio_ssd1306.SSD1306(
+        self.display = adafruit_displayio_ssd1306.SSD1306(
             displayio.I2CDisplay(self.i2c, device_address=self.device_address),
             width=width,
             height=height,
             rotation=rotation,
             brightness=self.brightness,
         )
 
+        return self.display
+
     def deinit(self):
         self.i2c.deinit()
 
+    def sleep(self):
+        self.display.sleep()
+
+    def wake(self):
+        self.display.wake()
+
 
 class SH1106(OledDisplayType):
     def __init__(
@@ -147,7 +170,7 @@ def __init__(
     def during_bootup(self, width, height, rotation):
         import adafruit_displayio_sh1106
 
-        return adafruit_displayio_sh1106.SH1106(
+        self.display = adafruit_displayio_sh1106.SH1106(
             displayio.FourWire(
                 self.spi,
                 command=self.command,
@@ -160,9 +183,17 @@ def during_bootup(self, width, height, rotation):
             rotation=rotation,
         )
 
+        return self.display
+
     def deinit(self):
         self.spi.deinit()
 
+    def sleep(self):
+        self.display.sleep()
+
+    def wake(self):
+        self.display.wake()
+
 
 class Oled(Extension):
     def __init__(
@@ -316,7 +347,7 @@ def dim(self):
                 and ticks_diff(ticks_ms(), self.timer_start)
                 > self.powersave_off_time_ms
             ):
-                self.display.sleep()
+                self.display_type.sleep()
 
             elif (
                 self.powersave_dim_time_ms
@@ -327,13 +358,13 @@ def dim(self):
 
             else:
                 self.display.brightness = self.brightness
-                self.display.wake()
+                self.display_type.wake()
 
         elif (
             self.off_time_ms
             and ticks_diff(ticks_ms(), self.timer_start) > self.off_time_ms
         ):
-            self.display.sleep()
+            self.display_type.sleep()
 
         elif (
             self.dim_time_ms
@@ -343,4 +374,4 @@ def dim(self):
 
         else:
             self.display.brightness = self.brightness
-            self.display.wake()
+            self.display_type.wake()