portal/examples/spi_panel_test.py

#!/usr/bin/env python3
"""Step-by-step drive test for a Pi-connected panel over SPI."""

from __future__ import annotations

import argparse
import os
import sys
import time

from leds.array_config import (
    MATRIX_BRIGHTNESS,
    MATRIX_SPI_BUS,
    MATRIX_SPI_DEVICE,
    matrix_pixel_index,
    panel_layout,
    panel_pixel_count,
)
from leds.backends.spi import spidev_bufsize_hint
from leds.config import panel_strip_config
from leds.matrix import LedMatrix


def _flash(matrix: LedMatrix, pause: float) -> None:
    print("Test: RGB flash")
    for name, color in (("red", (255, 0, 0)), ("green", (0, 255, 0)), ("blue", (0, 0, 255))):
        print(f"  {name}")
        matrix.fill(color)
        matrix.show()
        time.sleep(pause)
    matrix.clear()
    matrix.show()


def _first_last(matrix: LedMatrix, pixels: int, pause: float) -> None:
    print(f"Test: LED 0 red, LED {pixels - 1} green")
    width = matrix.width
    matrix.fill((0, 0, 0))
    matrix.set_pixel(0, 0, (255, 0, 0))
    last = pixels - 1
    matrix.set_pixel(last % width, last // width, (0, 255, 0))
    matrix.show()
    time.sleep(pause)
    matrix.clear()
    matrix.show()


def _rows(matrix: LedMatrix, width: int, height: int, pause: float) -> None:
    print("Test: one row at a time (white)")
    for y in range(height):
        matrix.fill((0, 0, 0))
        for x in range(width):
            matrix[x, y] = (255, 255, 255)
        print(f"  row {y}")
        matrix.show()
        time.sleep(pause)
    matrix.clear()
    matrix.show()


def _corners(matrix: LedMatrix, width: int, height: int, pause: float) -> None:
    print("Test: corners")
    corners = (
        ("top-left", 0, 0, (255, 0, 0)),
        ("top-right", width - 1, 0, (0, 255, 0)),
        ("bottom-left", 0, height - 1, (0, 0, 255)),
        ("bottom-right", width - 1, height - 1, (255, 255, 0)),
    )
    for label, x, y, color in corners:
        matrix.fill((0, 0, 0))
        matrix[x, y] = color
        idx = matrix_pixel_index(x, y, width, height)
        print(f"  {label} ({x},{y}) -> LED {idx}")
        matrix.show()
        time.sleep(pause)
    matrix.clear()
    matrix.show()


def _chase(matrix: LedMatrix, pixels: int, pause: float, steps: int) -> None:
    print(f"Test: chase ({steps} steps)")
    width = matrix.width
    for i in range(steps):
        matrix.fill((0, 0, 0))
        idx = i % pixels
        matrix[idx % width, idx // width] = (0, 255, 255)
        matrix.show()
        time.sleep(pause)
    matrix.clear()
    matrix.show()


def main() -> int:
    parser = argparse.ArgumentParser(description="SPI drive test for Pi-connected panel")
    parser.add_argument("--panel-index", type=int, default=2, choices=range(5), metavar="N")
    parser.add_argument("--brightness", type=float, default=MATRIX_BRIGHTNESS)
    parser.add_argument(
        "--spi-bus",
        type=int,
        default=int(os.environ.get("PORTAL_SPI_BUS", MATRIX_SPI_BUS)),
    )
    parser.add_argument("--spi-device", type=int, default=MATRIX_SPI_DEVICE)
    parser.add_argument("--spi-mhz", type=float, default=4.2)
    parser.add_argument("--wire-order", default=None)
    parser.add_argument("--pause", type=float, default=0.8)
    parser.add_argument(
        "--test",
        choices=("all", "flash", "first-last", "rows", "corners", "chase"),
        default="all",
    )
    parser.add_argument("--chase-steps", type=int, default=None)
    args = parser.parse_args()

    width, height = panel_layout(args.panel_index)
    pixels = panel_pixel_count(args.panel_index)
    hint = spidev_bufsize_hint(pixels)
    if hint:
        print(f"Warning: {hint}\n", file=sys.stderr)

    cfg = panel_strip_config(
        args.panel_index,
        pixels,
        spi_max_speed_hz=max(1, int(args.spi_mhz * 1_000_000)),
        wire_order=args.wire_order,
        brightness=args.brightness,
    )
    chase_steps = args.chase_steps if args.chase_steps is not None else pixels

    print(f"SPI panel test: panel {args.panel_index}, {width}×{height} ({pixels} LEDs)")

    try:
        with LedMatrix(width, height, config0=cfg) as matrix:
            if args.test in ("all", "flash"):
                _flash(matrix, args.pause)
            if args.test in ("all", "first-last"):
                _first_last(matrix, pixels, args.pause)
            if args.test in ("all", "corners"):
                _corners(matrix, width, height, args.pause)
            if args.test in ("all", "rows"):
                _rows(matrix, width, height, max(args.pause * 0.6, 0.15))
            if args.test in ("all", "chase"):
                _chase(matrix, pixels, max(args.pause * 0.25, 0.05), chase_steps)
    except KeyboardInterrupt:
        print("\nStopped.")
    except OSError as exc:
        print(f"SPI error: {exc}", file=sys.stderr)
        return 1
    print("Done.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())