ssd1306plus.py

# ssd1306plus - extended I2C and SPI SSD1306 oled driver
# v1.4.0 - True display rotation
# v1.3.0 - More efficient cropping
# v1.2.0 - GIF cropping, garbage collection
# v1.1.0 - added GIF support
# v1.0.0 - initial

from micropython import const
import framebuf, gc


# register definitions
SET_CONTRAST = const(0x81)
SET_ENTIRE_ON = const(0xA4)
SET_NORM_INV = const(0xA6)
SET_DISP = const(0xAE)
SET_MEM_ADDR = const(0x20)
SET_COL_ADDR = const(0x21)
SET_PAGE_ADDR = const(0x22)
SET_DISP_START_LINE = const(0x40)
SET_SEG_REMAP = const(0xA0)
SET_MUX_RATIO = const(0xA8)
SET_IREF_SELECT = const(0xAD)
SET_COM_OUT_DIR = const(0xC0)
SET_DISP_OFFSET = const(0xD3)
SET_COM_PIN_CFG = const(0xDA)
SET_DISP_CLK_DIV = const(0xD5)
SET_PRECHARGE = const(0xD9)
SET_VCOM_DESEL = const(0xDB)
SET_CHARGE_PUMP = const(0x8D)

# NOTE: We no longer subclass FrameBuffer because MicroPython's framebuf
# doesn't support re-calling FrameBuffer.__init__ on an existing instance.
class SSD1306:
    def __init__(self, width, height, external_vcc):
        # Physical geometry (panel RAM we send)
        self._phys_w = width
        self._phys_h = height
        self._phys_pages = self._phys_h // 8
        self._phys_buf = bytearray(self._phys_pages * self._phys_w)

        self.external_vcc = external_vcc
        self._rot = 0  # 0, 90, 180, 270

        # Physical framebuffer (always width x height)
        self._fb_phys = framebuf.FrameBuffer(self._phys_buf, self._phys_w, self._phys_h, framebuf.MONO_VLSB)

        # Logical framebuffer for 90/270 (swapped)
        self._log_w = self._phys_h
        self._log_h = self._phys_w
        self._log_pages = self._log_h // 8
        self._log_buf = bytearray(self._log_pages * self._log_w)
        self._fb_log = framebuf.FrameBuffer(self._log_buf, self._log_w, self._log_h, framebuf.MONO_VLSB)

        # Active drawing target (starts unrotated)
        self.width = self._phys_w
        self.height = self._phys_h
        self.pages = self._phys_pages
        self.buffer = self._phys_buf
        self._fb = self._fb_phys

        self.init_display()

    # ---- Forward FrameBuffer API we rely on (plus code uses these) ----
    def fill(self, c):
        return self._fb.fill(c)

    def pixel(self, x, y, c=None):
        if c is None:
            return self._fb.pixel(x, y)
        return self._fb.pixel(x, y, c)

    def fill_rect(self, x, y, w, h, c):
        return self._fb.fill_rect(x, y, w, h, c)

    def rect(self, x, y, w, h, c):
        return self._fb.rect(x, y, w, h, c)

    def hline(self, x, y, w, c):
        return self._fb.hline(x, y, w, c)

    def vline(self, x, y, h, c):
        return self._fb.vline(x, y, h, c)

    def line(self, x1, y1, x2, y2, c):
        return self._fb.line(x1, y1, x2, y2, c)

    def text(self, font, s, x, y, c=1):
        """
        Keep compatibility with your current usage:
          tft.text(vga2_16x16, "hi", x, y, col)
        On SSD1306plus you might call either FrameBuffer.text(str, x, y)
        or a custom font system. Your current code uses framebuf text in scaled(),
        and gif() uses pixel() directly.
        """
        # If caller passed a font module (vga1_8x8 etc) like on your GC9A01 code,
        # framebuf can't use that. For SSD1306, existing projects typically call:
        #   oled.text("hi", x, y, 1)
        # Your file currently calls temp_fb.text(text, 0, 0) in scaled().
        #
        # So: support both signatures:
        #   text("hi", x, y, 1)
        #   text(font, "hi", x, y, 1)  -> we ignore font and use builtin
        if isinstance(font, str):
            # signature text(s, x, y, c=1)
            s2 = font
            x2 = s
            y2 = x
            c2 = y if y is not None else 1
            return self._fb.text(s2, x2, y2, c2)

        # signature text(font, s, x, y, c=1) -> ignore font
        return self._fb.text(s, x, y, c)

    def blit(self, fbuf, x, y, key=-1):
        return self._fb.blit(fbuf, x, y, key)

    # ---- Bit helpers for software rotation (VLSB) ----
    def _vlsb_get(self, buf, w, x, y):
        i = x + (y >> 3) * w
        return (buf[i] >> (y & 7)) & 1

    def _vlsb_set(self, buf, w, x, y, v):
        i = x + (y >> 3) * w
        m = 1 << (y & 7)
        if v:
            buf[i] |= m
        else:
            buf[i] &= ~m

    def init_display(self):
        for cmd in (
            SET_DISP,  # display off
            # address setting
            SET_MEM_ADDR,
            0x00,  # horizontal
            # resolution and layout
            SET_DISP_START_LINE,  # start at line 0
            SET_SEG_REMAP | 0x01,  # column addr 127 mapped to SEG0
            SET_MUX_RATIO,
            self._phys_h - 1,
            SET_COM_OUT_DIR | 0x08,  # scan from COM[N] to COM0
            SET_DISP_OFFSET,
            0x00,
            SET_COM_PIN_CFG,
            0x02 if self._phys_w > 2 * self._phys_h else 0x12,
            # timing and driving scheme
            SET_DISP_CLK_DIV,
            0x80,
            SET_PRECHARGE,
            0x22 if self.external_vcc else 0xF1,
            SET_VCOM_DESEL,
            0x30,  # 0.83*Vcc
            # display
            SET_CONTRAST,
            0xFF,  # maximum
            SET_ENTIRE_ON,  # output follows RAM contents
            SET_NORM_INV,  # not inverted
            SET_IREF_SELECT,
            0x30,  # enable internal IREF during display on
            # charge pump
            SET_CHARGE_PUMP,
            0x10 if self.external_vcc else 0x14,
            SET_DISP | 0x01,  # display on
        ):  # on
            self.write_cmd(cmd)
        self.fill(0)
        self.show()

    def poweroff(self):
        self.write_cmd(SET_DISP)

    def poweron(self):
        self.write_cmd(SET_DISP | 0x01)

    def contrast(self, contrast):
        self.write_cmd(SET_CONTRAST)
        self.write_cmd(contrast)

    def invert(self, invert):
        self.write_cmd(SET_NORM_INV | (invert & 1))

    def rotate(self, rotate):
        """
        Backwards compatible with your old behavior:
          rotate(0) -> normal
          rotate(1) -> 180 flip (hardware)

        New:
          rotate(90), rotate(180), rotate(270)
        """
        if rotate in (0, 1):
            deg = 180 if rotate else 0
        else:
            deg = rotate

        deg = deg % 360
        if deg not in (0, 90, 180, 270):
            deg = 0
        self._rot = deg

        if deg in (0, 180):
            self.width = self._phys_w
            self.height = self._phys_h
            self.pages = self._phys_pages
            self.buffer = self._phys_buf
            self._fb = self._fb_phys

            flip = 1 if deg == 180 else 0
            self.write_cmd(SET_COM_OUT_DIR | (flip << 3))
            self.write_cmd(SET_SEG_REMAP | flip)
        else:
            self.width = self._log_w
            self.height = self._log_h
            self.pages = self._log_pages
            self.buffer = self._log_buf
            self._fb = self._fb_log

            # keep panel normal; software rotate in show()
            self.write_cmd(SET_COM_OUT_DIR | (0 << 3))
            self.write_cmd(SET_SEG_REMAP | 0)

        self.fill(0)
        self.show()

    def gif(self, filename, x=0, y=0, loop=1, delay_ms=None, clear=False, crop=None):
        """
        - filename: path to .gif file on the filesystem
        - x, y: offset of the GIF's logical screen on the display
        - loop: number of times to loop; use -1 for infinite
        - delay_ms: override frame delay (ms). If None, use GIF's own delay.
        - clear: if True, clear the frame area before drawing each frame
        - Supports non-interlaced GIFs only.
        - Blocking: does not return until all loops are done.
        """
        import time

        # Read entire GIF file
        with open(filename, "rb") as f:
            data = f.read()

        if len(data) < 13:
            return

        # Header: "GIF87a" or "GIF89a"
        if not (data.startswith(b"GIF87a") or data.startswith(b"GIF89a")):
            return

        # Logical Screen Descriptor
        ls_width  = data[6] | (data[7] << 8)
        ls_height = data[8] | (data[9] << 8)
        packed    = data[10]
        bg_color_index = data[11]
        # pixel_aspect = data[12]  # not used

        # Global Color Table
        gct_flag = (packed & 0x80) != 0
        gct_size = 0
        global_color_table = None

        p = 13  # position after header + LSD
        if gct_flag:
            gct_size = 1 << ((packed & 0x07) + 1)
            gct_bytes = 3 * gct_size
            global_color_table = data[p:p + gct_bytes]
            p += gct_bytes

        # Helper: read GIF sub-blocks into one bytes object
        def _read_subblocks(pos):
            img_data = bytearray()
            length = len(data)
            while pos < length:
                block_len = data[pos]
                pos += 1
                if block_len == 0:
                    break
                img_data.extend(data[pos:pos + block_len])
                pos += block_len
            return bytes(img_data), pos

        # Helper: LZW decode the image data into colour indices
        def _lzw_decode(img_bytes, min_code_size, expected_pixels):
            # Simple GIF LZW decoder, up to 12-bit codes
            if min_code_size < 2 or min_code_size > 8:
                # Fallback: just truncate raw bytes
                return list(img_bytes[:expected_pixels])

            clear_code = 1 << min_code_size
            end_code = clear_code + 1
            code_size = min_code_size + 1
            max_code_size = 12

            # Dictionary: code -> [indices...]
            dictionary = [[i] for i in range(clear_code)] + [None, None]

            bit_pos = 0
            data_len = len(img_bytes)
            output = []
            prev = None

            def _next_code(bit_pos, code_size):
                byte_pos = bit_pos >> 3
                if byte_pos >= data_len:
                    return None, bit_pos
                raw = 0
                bits_read = 0
                shift = 0
                while bits_read < code_size and byte_pos < data_len:
                    b = img_bytes[byte_pos]
                    available = 8 - (bit_pos & 7)
                    take = code_size - bits_read
                    if take > available:
                        take = available
                    mask = (1 << take) - 1
                    v = (b >> (bit_pos & 7)) & mask
                    raw |= v << shift
                    bits_read += take
                    shift += take
                    bit_pos += take
                    byte_pos = bit_pos >> 3
                if bits_read != code_size:
                    return None, bit_pos
                return raw, bit_pos

            # Initial code
            code, bit_pos = _next_code(bit_pos, code_size)
            if code is None:
                return output

            if code == clear_code:
                dictionary = [[i] for i in range(clear_code)] + [None, None]
                code_size = min_code_size + 1
                code, bit_pos = _next_code(bit_pos, code_size)
                if code is None or code == end_code:
                    return output

            if code == end_code:
                return output
            if code >= len(dictionary) or dictionary[code] is None:
                return output

            prev = dictionary[code][:]
            output.extend(prev)

            while len(output) < expected_pixels:
                code, bit_pos = _next_code(bit_pos, code_size)
                if code is None:
                    break

                if code == clear_code:
                    dictionary = [[i] for i in range(clear_code)] + [None, None]
                    code_size = min_code_size + 1
                    code, bit_pos = _next_code(bit_pos, code_size)
                    if code is None or code == end_code:
                        break
                    if code >= len(dictionary) or dictionary[code] is None:
                        cur = []
                    else:
                        cur = dictionary[code][:]
                    output.extend(cur)
                    prev = cur[:]
                    continue

                if code == end_code:
                    break

                if code < len(dictionary) and dictionary[code] is not None:
                    cur = dictionary[code][:]
                elif code == len(dictionary):
                    # KwKwK case
                    cur = prev[:] + [prev[0]]
                else:
                    # Invalid code
                    break

                output.extend(cur)
                dictionary.append(prev[:] + [cur[0]])

                if len(dictionary) == (1 << code_size) and code_size < max_code_size:
                    code_size += 1

                prev = cur[:]

            return output[:expected_pixels]

        # Animation loop
        loops_done = 0
        infinite = loop < 0
        start_pos = p  # where the blocks start

        while infinite or loops_done < loop:
            p = start_pos
            transparency_index = None
            frame_delay_ms = 0

            while p < len(data):
                block_type = data[p]
                p += 1

                # Trailer: end of GIF
                if block_type == 0x3B:
                    break

                # Extension block
                if block_type == 0x21:
                    label = data[p]
                    p += 1

                    # Graphics Control Extension (gives us delay + transparency)
                    if label == 0xF9:
                        block_size = data[p]
                        p += 1
                        if block_size == 4:
                            packed_fields = data[p]
                            p += 1
                            delay_lo = data[p]
                            delay_hi = data[p + 1]
                            p += 2
                            trans_index = data[p]
                            p += 1
                            terminator = data[p]
                            p += 1

                            frame_delay_ms = (delay_hi << 8 | delay_lo) * 10
                            if frame_delay_ms <= 0:
                                frame_delay_ms = 50  # sensible default

                            if packed_fields & 0x01:
                                transparency_index = trans_index
                        else:
                            sub_len = data[p]
                            p += 1 + sub_len
                            while sub_len:
                                sub_len = data[p]
                                p += 1 + sub_len

                    else:
                        sub_len = data[p]
                        p += 1
                        while sub_len:
                            p += sub_len
                            sub_len = data[p]
                            p += 1

                    continue

                # Image Descriptor
                if block_type == 0x2C:
                    left = data[p] | (data[p + 1] << 8)
                    top = data[p + 2] | (data[p + 3] << 8)
                    width = data[p + 4] | (data[p + 5] << 8)
                    height = data[p + 6] | (data[p + 7] << 8)
                    packed_fields = data[p + 8]
                    p += 9

                    lct_flag = (packed_fields & 0x80) != 0
                    interlace = (packed_fields & 0x40) != 0

                    if interlace:
                        if lct_flag:
                            lct_size = 1 << ((packed_fields & 0x07) + 1)
                            p += 3 * lct_size
                        _, p = _read_subblocks(p + 1)
                        continue

                    if lct_flag:
                        lct_size = 1 << ((packed_fields & 0x07) + 1)
                        p += 3 * lct_size

                    lzw_min_code_size = data[p]
                    p += 1

                    img_bytes, p = _read_subblocks(p)
                    expected_pixels = width * height
                    indices = _lzw_decode(img_bytes, lzw_min_code_size, expected_pixels)

                    if clear:
                        self.fill_rect(x + left, y + top, width, height, 0)

                    pos = 0
                    bg_idx = bg_color_index

                    dst_x0 = x + left
                    dst_y0 = y + top

                    if crop is not None:
                        cx0, cy0, cx1, cy1 = crop
                    else:
                        cx0 = cy0 = cx1 = cy1 = 0

                    for yy in range(height):
                        y_pix = dst_y0 + yy
                        for xx in range(width):
                            if pos >= len(indices):
                                break
                            idx = indices[pos]
                            pos += 1

                            if (transparency_index is not None) and (idx == transparency_index):
                                continue

                            x_pix = dst_x0 + xx

                            if crop is not None:
                                if (x_pix < cx0) or (x_pix > cx1) or (y_pix < cy0) or (y_pix > cy1):
                                    continue

                            col = 0
                            if idx != bg_idx:
                                col = 1

                            self.pixel(x_pix, y_pix, col)

                    self.show()

                    d = delay_ms if (delay_ms is not None) else frame_delay_ms
                    if d <= 0:
                        d = 50
                    time.sleep_ms(d)

                    continue
            gc.collect()
            loops_done += 1

    def play_gif(self, filename, x=0, y=0, loop=1, delay_ms=None, clear=False, crop=None):
        return self.gif(filename, x, y, loop, delay_ms, clear, crop)

    def show(self):
        # 90/270: rotate logical buffer -> physical buffer
        if self._rot in (90, 270):
            pb = self._phys_buf
            for i in range(len(pb)):
                pb[i] = 0

            lw, lh = self._log_w, self._log_h
            pw, ph = self._phys_w, self._phys_h
            lb = self._log_buf

            if self._rot == 90:
                # phys_x = y, phys_y = ph-1-x
                for y in range(lh):
                    for x in range(lw):
                        if self._vlsb_get(lb, lw, x, y):
                            self._vlsb_set(pb, pw, y, (ph - 1 - x), 1)
            else:
                # 270: phys_x = pw-1-y, phys_y = x
                for y in range(lh):
                    for x in range(lw):
                        if self._vlsb_get(lb, lw, x, y):
                            self._vlsb_set(pb, pw, (pw - 1 - y), x, 1)

            out_buf = pb
            out_w = pw
            out_pages = self._phys_pages
        else:
            out_buf = self._phys_buf
            out_w = self._phys_w
            out_pages = self._phys_pages

        x0 = 0
        x1 = out_w - 1
        if out_w != 128:
            col_offset = (128 - out_w) // 2
            x0 += col_offset
            x1 += col_offset
        self.write_cmd(SET_COL_ADDR)
        self.write_cmd(x0)
        self.write_cmd(x1)
        self.write_cmd(SET_PAGE_ADDR)
        self.write_cmd(0)
        self.write_cmd(out_pages - 1)
        self.write_data(out_buf)


class SSD1306_I2C(SSD1306):
    def __init__(self, width, height, i2c, addr=0x3C, external_vcc=False):
        self.i2c = i2c
        self.addr = addr
        self.temp = bytearray(2)
        self.write_list = [b"\x40", None]  # Co=0, D/C#=1
        super().__init__(width, height, external_vcc)

    def write_cmd(self, cmd):
        self.temp[0] = 0x80  # Co=1, D/C#=0
        self.temp[1] = cmd
        self.i2c.writeto(self.addr, self.temp)

    def write_data(self, buf):
        self.write_list[1] = buf
        self.i2c.writevto(self.addr, self.write_list)

    def scaled(self, text, x, y, scale=2, colr=1):
        clr = 0
        temp_w = len(text) * 8
        temp_h = 8
        temp_buf = bytearray((temp_w // 8) * temp_h)
        temp_fb = framebuf.FrameBuffer(temp_buf, temp_w, temp_h, framebuf.MONO_HLSB)
        temp_fb.fill(0)
        temp_fb.text(text, 0, 0)

        for ix in range(temp_w):
            for iy in range(temp_h):
                if temp_fb.pixel(ix, iy):
                    x0 = x + ix * scale
                    y0 = y + iy * scale
                    for dx in range(scale):
                        for dy in range(scale):
                            self.pixel(x0 + dx, y0 + dy, colr)


class SSD1306_SPI(SSD1306):
    def __init__(self, width, height, spi, dc, res, cs, external_vcc=False):
        self.rate = 10 * 1024 * 1024
        dc.init(dc.OUT, value=0)
        res.init(res.OUT, value=0)
        cs.init(cs.OUT, value=1)
        self.spi = spi
        self.dc = dc
        self.res = res
        self.cs = cs
        import time

        self.res(1)
        time.sleep_ms(1)
        self.res(0)
        time.sleep_ms(10)
        self.res(1)
        super().__init__(width, height, external_vcc)

    def write_cmd(self, cmd):
        self.spi.init(baudrate=self.rate, polarity=0, phase=0)
        self.cs(1)
        self.dc(0)
        self.cs(0)
        self.spi.write(bytearray([cmd]))
        self.cs(1)

    def write_data(self, buf):
        self.spi.init(baudrate=self.rate, polarity=0, phase=0)
        self.cs(1)
        self.dc(1)
        self.cs(0)
        self.spi.write(buf)
        self.cs(1)

    def scaled(self, text, x, y, scale=2, colr=1):
        clr = 0
        temp_w = len(text) * 8
        temp_h = 8
        temp_buf = bytearray((temp_w // 8) * temp_h)
        temp_fb = framebuf.FrameBuffer(temp_buf, temp_w, temp_h, framebuf.MONO_HLSB)
        temp_fb.fill(0)
        temp_fb.text(text, 0, 0)

        for ix in range(temp_w):
            for iy in range(temp_h):
                if temp_fb.pixel(ix, iy):
                    x0 = x + ix * scale
                    y0 = y + iy * scale
                    for dx in range(scale):
                        for dy in range(scale):
                            self.pixel(x0 + dx, y0 + dy, colr)