OmniaLeds/omnia/leds.py

from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Union
from enum import Enum, auto

@dataclass
class LedState:
	r: Optional[int]
	g: Optional[int]
	b: Optional[int]
	brightness: Optional[int]
	autonomous: Optional[bool]


class LED:
	def __init__(self, alias):
		self.alias = alias
		self.name = unalias(alias)
		self.dir = Path('/sys/class/leds/'+self.name+'/')

	@property
	def state(self) -> LedState:
		with open(self.dir/'color') as f:
			r, g, b = [int(x) for x in f.read().strip().split()]
		with open(self.dir/'brightness') as f:
			brightness = int(f.read.strip())
		with open(self.dir/'autonomous') as f:
			autonomous = bool(f.read.strip())
		return LedState(r=r, g=g, b=b, brightness=brightness, autonomous=autonomous)

	def set(self, state: LedState):
		# TODO: sanity check: when autonomous, no color or warn
		
		# None means do not change value
		old_state = self.state
		for attr in ['r', 'g', 'b', 'brightness', 'autonomous']:
			if getattr(state, attr) == None:
				old_value = getattr(old_state, attr)
				setattr(state, attr, old_value)

		# Set it
		with open(self.dir/'color', 'w') as f:
			f.write(f'{state.r} {state.g} {state.b}\n')
		with open(self.dir/'brightness', 'w') as f:
			f.write(str(state.brightness)+'\n')
		with open(self.dir/'autonomous', 'w') as f:
			f.write(str(int(state.autonomous))+'\n')

def unalias(s):
	return 'omnia-led:' + s

def alias(s):
	return s.removeprefix('omnia-led:')

class PatternFitting(Enum):
	Repeat = auto()
	Nearest = auto()
	Linear = auto()

class LedStrip:
	"Represents a line of LEDs"

	def __init__(self, leds: Union[list[str], list[LED]]):
		if all(isinstance(x, str) for x in leds):
			leds = [LED(x) for x in leds]
		self.leds = leds

	def __len__(self):
		return len(self.leds)
	def __iter__(self):
		return self.leds.__iter__()

	@property
	def pattern(self) -> list[LedState]:
		return [x.state for x in self.leds]
	def setpattern(self, pat: list[LedState]):
		if len(pat) != len(self):
			raise ValueError('The pattern must be a one-to-one matching to LEDs')
		for led, st in zip(self.leds, pat):
			led.set(st)

	def apply(self, pat: list[LedState], fit: PatternFitting):
		"""Sets the strip to given pattern, which will be fit using selected algorithm.

		Unlike setpattern, this first / always modifies the pattern to somehow fit to the given strip
		"""
		patl = len(pat)
		strl = len(self)

		if False: pass
		elif fit == fit.Repeat:
			# Repeat and/or chop
			new_pat = (pat * ((strl // patl) + 1))[:strl]
		elif fit == fit.Nearest:
			quantum = patl/strl
			indices = [round(quantum * i) for i in range(strl)]
			new_pat = [pat[i] for i in indices]
		elif fit == fit.Linear:
			# Too complex
			raise NotImplementedError('Please submit a pull request :-)')

		self.setpattern(new_pat)

omnia_led_order = [
	'user1',
	'user2',
	'lan0',
	'lan1',
	'lan2',
	'lan3',
	'lan4',
	'wan',
	'pci1',
	'pci2',
	'pci3',
	'power',
	]

OmniaStrip = LedStrip(omnia_led_order)