Host: Made rendering pipeline a bit more flexible

8 files changed, 125 insertions, 96 deletions

diff --git a/host/bdf.py b/host/bdf.py
index b386426..a8fc8f3 100644
--- a/host/bdf.py
+++ b/host/bdf.py
@@ -27,7 +27,6 @@ def printframe(fb):
 		numpy.copyto(dbuf[2::4], ip[2::3+rgba])
 		lib.console_render_buffer(dbuf.ctypes.data_as(POINTER(c_uint8)), config.display_width, config.display_height)
 
-
 class Font:
 	def __init__(self, fontfile='unifont.bdf'):
 		self.font = lib.read_bdf_file(fontfile)
@@ -36,16 +35,16 @@ class Font:
 		self.cbuf = create_string_buffer(config.frame_size*sizeof(COLOR))
 		self.cbuflock = threading.Lock()
 
-	def compute_text_bounds(text):
+	def compute_text_bounds(self, text):
 		textbytes = text.encode()
 		textw, texth = c_size_t(0), c_size_t(0)
-		res = lib.framebuffer_get_text_bounds(textbytes, unifont, byref(textw), byref(texth))
+		res = lib.framebuffer_get_text_bounds(textbytes, self.font, byref(textw), byref(texth))
 		if res:
 			raise ValueError('Invalid text')
 		return textw.value, texth.value
 
-	def render_text(text, offset):
-		with cbuflock:
+	def render_text(self, text, offset):
+		with self.cbuflock:
 			textbytes = bytes(str(text), 'UTF-8')
 			res = lib.framebuffer_render_text(textbytes, self.font, self.cbuf,
 					config.display_width, config.display_height, offset)
@@ -53,4 +52,4 @@ class Font:
 				raise ValueError('Invalid text')
 			return self.cbuf
 
-
+unifont = Font()
diff --git a/host/config.py b/host/config.py
index c6bd42e..02a047c 100644
--- a/host/config.py
+++ b/host/config.py
@@ -41,3 +41,6 @@ udp_port = tcp_port = 1337
 # Forward addr/port
 crap_fw_addr, crap_fw_port = '127.0.0.1', 1338
 
+# USB Serial number of matelight to control as byte string (None for first matelight connected)
+ml_usb_serial_match = None
+
diff --git a/host/crap.py b/host/crap.py
index db9f7ee..80fb647 100644
--- a/host/crap.py
+++ b/host/crap.py
@@ -4,17 +4,20 @@ import struct
 import zlib
 import io
 from time import time
+import numpy
 
 import config
 
 class CRAPClient:
 	def __init__(self, ip='127.0.0.1', port=1337):
 		self.ip, self.port = ip, port
-		self.sock = socket.Socket(socket.AF_INET, socket.SOCK_DGRAM)
+		self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
 		self.close = self.sock.close
 
 	def sendframe(self, frame):
-		self.sock.sendto(frame, (self.ip, self.port))
+		fb = numpy.frombuffer(frame, dtype=numpy.uint8)
+		fb.shape = config.frame_size, len(frame)/config.frame_size
+		self.sock.sendto(fb[:,:3].tobytes(), (self.ip, self.port))
 
 
 def _timestamped_recv(sock):
@@ -23,6 +26,8 @@ def _timestamped_recv(sock):
 			data, addr = sock.recvfrom(config.frame_size*3+4)
 		except io.BlockingIOError as e:
 			raise StopIteration()
+		except socket.timeout:
+			raise StopIteration()
 		else:
 			yield time(), data, addr
 
@@ -34,6 +39,7 @@ class CRAPServer:
 		self.sock.setblocking(blocking)
 		self.sock.bind((ip, port))
 
+		self.blocking = blocking
 		self.current_client = None
 		self.last_timestamp = 0
 		self.begin_timestamp = 0
@@ -45,13 +51,14 @@ class CRAPServer:
 	def __iter__(self):
 		for timestamp, data, (addr, sport) in _timestamped_recv(self.sock):
 			if data is None:
-				yield None
+				yield None, None
 
-			if timestamp - self.last_timestamp > config.udp_timeout\
+			if timestamp - self.last_timestamp >= config.udp_timeout\
 					or timestamp - self.begin_timestamp > config.udp_switch_interval:
 				self.current_client = addr
 				self.begin_timestamp = timestamp
 				self.log('\x1B[91mAccepting UDP data from\x1B[0m', addr)
+				self.sock.settimeout(config.udp_timeout)
 
 			if addr == self.current_client:
 				if len(data) == config.frame_size*3+4:
@@ -63,6 +70,8 @@ class CRAPServer:
 				elif len(data) != config.frame_size*3:
 					self.log('Error receiving UDP frame: Invalid frame size: {}'.format(len(data)))
 				self.last_timestamp = timestamp
-				yield data
+				yield 'udp:'+addr, data
+		self.current_client = None
+		self.sock.settimeout(None if self.blocking else 0)
 
 
diff --git a/host/matelight.py b/host/matelight.py
index 9966aee..886e0f1 100644
--- a/host/matelight.py
+++ b/host/matelight.py
@@ -3,28 +3,36 @@ from itertools import product
 from ctypes import c_size_t, c_uint8, c_void_p, c_float, CDLL, Structure, POINTER
 import numpy as np
 import time
+import atexit
 
-from config import *
+import config
 
 ml = CDLL('./libml.so')
 ml.matelight_open.restype = c_void_p
 
 if ml.matelight_usb_init():
 	raise OSError('Cannot initialize USB library')
-matelights = ml.matelight_open()
-if matelights is None:
-	raise ImportError('Cannot open any Mate Light devices')
-
-dbuf = np.zeros(DISPLAY_WIDTH*DISPLAY_HEIGHT*4, dtype=np.uint8)
-def sendframe(framedata):
-	""" Send a frame to the display
-
-	The argument contains a h * w array of 3-tuples of (r, g, b)-data or 4-tuples of (r, g, b, a)-data where the a
-	channel is ignored.
-	"""
-	# just use the first Mate Light available
-	rgba = len(framedata) == DISPLAY_WIDTH*DISPLAY_HEIGHT*4
-	global dbuf
-	np.copyto(dbuf[:640*(3+rgba)], np.frombuffer(framedata, dtype=np.uint8))
-	ml.matelight_send_frame(matelights, dbuf.ctypes.data_as(POINTER(c_uint8)), c_size_t(CRATES_X), c_size_t(CRATES_Y), c_float(BRIGHTNESS), rgba)
+
+atexit.register(ml.matelight_usb_destroy)
+
+class Matelight:
+	def __init__(self, match_serial=None):
+		""" Open the matelight matching the USB serial number given as a bytes object. If match_serial is None, open the
+		first matelight """
+		self.handle = ml.matelight_open(match_serial)
+		self.dbuf = np.zeros(config.frame_size*4, dtype=np.uint8)
+		if self.handle is None:
+			raise ValueError('Cannot find requested matelight.')
+
+	def sendframe(self, framedata):
+		""" Send a frame to the display
+
+		The argument contains a h * w array of 3-tuples of (r, g, b)-data or 4-tuples of (r, g, b, a)-data where the a
+		channel is ignored.
+		"""
+		# just use the first Mate Light available
+		rgba = len(framedata) == config.frame_size*4
+		np.copyto(self.dbuf[:640*(3+rgba)], np.frombuffer(framedata, dtype=np.uint8))
+		ml.matelight_send_frame(self.handle, self.dbuf.ctypes.data_as(POINTER(c_uint8)), c_size_t(config.crates_x),
+				c_size_t(config.crates_y), c_float(config.brightness), rgba)
 
diff --git a/host/server.py b/host/server.py
index 3ce1bdc..f83e2cd 100755
--- a/host/server.py
+++ b/host/server.py
@@ -1,12 +1,14 @@
 #!/usr/bin/env python
 
 import socket
-from time import strftime
+import time
 import itertools
 import sys
 from contextlib import suppress
+import asyncio
+import threading
 
-from config import *
+import config
 
 import matelight
 import bdf
@@ -14,48 +16,48 @@ import crap
 
 
 def log(*args):
-	print(strftime('\x1B[93m[%m-%d %H:%M:%S]\x1B[0m'), ' '.join(str(arg) for arg in args), '\x1B[0m')
+	print(time.strftime('\x1B[93m[%m-%d %H:%M:%S]\x1B[0m'), ' '.join(str(arg) for arg in args), '\x1B[0m')
 	sys.stdout.flush()
 
 class TextRenderer:
-	def __init__(self, text):
+	def __init__(self, text, title='default', font=bdf.unifont):
 		self.text = text
-		self.width, _ = unifont.compute_text_bounds(text)
+		self.font = font
+		(self.width, _), _testrender  = font.compute_text_bounds(text), font.render_text(text, 0)
+		self.title = title
 
 	def __iter__(self):
-		for i in range(-DISPLAY_WIDTH, self.width):
-			yield render_text(self.text, i)
+		for i in range(-config.display_width, self.width):
+			yield self.title, self.font.render_text(self.text, i)
+			time.sleep(0.05)
 
 class MatelightTCPServer:
-	def __init__(self, port, ip):
-		self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-		self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-		self.sock.setblocking(blocking)
-		self.sock.bind((ip, port))
-		self.conns = set()
+	def __init__(self, ip, port, loop):
+		coro = asyncio.start_server(self.handle_conn, ip, port, loop=loop)
+		server = loop.run_until_complete(coro)
 		self.renderqueue = []
+		self.close = server.close
 
 	def __iter__(self):
 		q, self.renderqueue = self.renderqueue, []
-		for frame in itertools.chain(*q):
-			yield frame
+		for title, frame in itertools.chain(*q):
+			yield title, frame
 
-	def handle_connections(self):
-		for conn in self.conns:
+	async def handle_conn(self, reader, writer):
+		line = (await reader.read(1024)).decode('UTF-8').strip()
+		if len(line) > 140: # Unicode string length, *not* byte length of encoded UTF-8
+			writer.write(b'TOO MUCH INFORMATION!\n')
+		else:
+			addr,*rest = writer.get_extra_info('peername')
+			log('\x1B[95mText from\x1B[0m {}: {}\x1B[0m'.format(addr, line))
 			try:
-				line = conn.recv(1024).decode('UTF-8').strip()
-				if len(data) > 140: # Unicode string length, *not* byte length of encoded UTF-8
-					conn.sendall(b'TOO MUCH INFORMATION!\n')
-				else:
-					log('\x1B[95mText from\x1B[0m {}: {}\x1B[0m'.format(addr, data))
-					renderqueue.append(TextRenderer(data))
-					conn.sendall(b'KTHXBYE!\n')
-			except socket.error, e:
-				if err == errno.EAGAIN or err == errno.EWOULDBLOCK:
-					continue
-			with suppress(socket.error):
-				conn.close()
-			self.conns.remove(conn)
+				self.renderqueue.append(TextRenderer(line, title='tcp:'+addr))
+			except:
+				writer.write(b'STAHPTROLLINK?\n')
+			else:
+				writer.write(b'KTHXBYE!\n')
+		await writer.drain()
+		writer.close()
 
 def _fallbackiter(it, fallback):
 	for fel in fallback:
@@ -63,19 +65,32 @@ def _fallbackiter(it, fallback):
 			yield el
 		yield fel
 
+def defaulttexts(filename='default.lines'):
+	with open(filename) as f:
+		return itertools.chain.from_iterable(( TextRenderer(l[:-1].replace('\\x1B', '\x1B')) for l in f.readlines() ))
+
 if __name__ == '__main__':
-	tcp_server = MatelightTCPServer(config.tcp_addr, config.tcp_port)
+	try:
+		ml = matelight.Matelight(config.ml_usb_serial_match)
+	except ValueError as e:
+		print(e, 'Starting in headless mode.', file=sys.stderr)
+		ml = None
+
+	loop = asyncio.get_event_loop()
+
+	tcp_server = MatelightTCPServer(config.tcp_addr, config.tcp_port, loop)
 	udp_server = crap.CRAPServer(config.udp_addr, config.udp_port)
 	forwarder  = crap.CRAPClient(config.crap_fw_addr, config.crap_fw_port) if config.crap_fw_addr is not None else None
 
-	def defaulttexts(filename='default.lines'):
-		with open(filename) as f:
-			return itertools.chain.from_iterable(( TextRenderer(l[:-1].replace('\\x1B', '\x1B')) for l in f.readlines() ))
+	async_thr = threading.Thread(target=loop.run_forever)
+	async_thr.daemon = True
+	async_thr.start()
 
 	with suppress(KeyboardInterrupt):
-		for renderer in _fallbackiter(tcp_server, defaulttexts()):
-			for frame in _fallbackiter(udp_server, renderer):
-				matelight.sendframe(frame)
+		while True:
+			for title, frame in _fallbackiter(udp_server, _fallbackiter(tcp_server, defaulttexts())):
+				if ml:
+					ml.sendframe(frame)
 				if forwarder:
 					forwarder.sendframe(frame)
 
diff --git a/host/usb.c b/host/usb.c
index 5e696d5..18b38a8 100644
--- a/host/usb.c
+++ b/host/usb.c
@@ -32,10 +32,12 @@ static int matelight_cmp_str_desc(libusb_device_handle *dev, uint8_t index, char
 	return strcmp(data, value) == 0;
 }
 
-matelight_handle *matelight_open(){
+matelight_handle *matelight_open(char *match_serial){
 	matelight_handle *out = NULL;
+    libusb_device_handle *handle = NULL;
 	libusb_device** device_list;
 	struct libusb_device_descriptor desc;
+
 	ssize_t res = libusb_get_device_list(NULL, &device_list);
 	if(res == 0){
 		fprintf(stderr, "LIBML: Cannot find any connected matelight\n");
@@ -44,56 +46,49 @@ matelight_handle *matelight_open(){
 		fprintf(stderr, "LIBML: Error enumerating connected USB devices\n");
 		goto error;
 	}else{
-		out = calloc(res+1, sizeof(matelight_handle));
+		out = calloc(1, sizeof(matelight_handle));
 		if(!out){
 			fprintf(stderr, "LIBML: Cannot allocate memory\n");
 			goto error;
 		}
-		memset(out, 0, (res+1)*sizeof(matelight_handle));
-		unsigned int found = 0;
 		for(ssize_t i=0; i<res; i++){
 			libusb_get_device_descriptor(device_list[i], &desc);
 			if(desc.idVendor == MATELIGHT_VID && desc.idProduct == MATELIGHT_PID){
-				libusb_device_handle *handle;
-				if(libusb_open(device_list[i], &(handle))){
+				if(libusb_open(device_list[i], &handle)){
 					fprintf(stderr, "LIBML: Cannot open Mate Light USB device\n");
 					goto error;
 				}
-				out[found].handle = handle;
+				out->handle = handle;
 				if(matelight_cmp_str_desc(handle, desc.iManufacturer, "Gold & Apple"))
-				if(matelight_cmp_str_desc(handle, desc.iProduct, "Mate Light")){
+				if(matelight_cmp_str_desc(handle, desc.iProduct, "Mate Light"))
+				if(!match_serial || matelight_cmp_str_desc(handle, desc.iSerialNumber, match_serial)){
 #define BUF_SIZE 256
-					char *serial = malloc(BUF_SIZE);
-					if(!serial){
-						fprintf(stderr, "LIBML: Cannot allocate memory\n");
+					out->serial = malloc(BUF_SIZE);
+					if(!out->serial){ fprintf(stderr, "LIBML: Cannot allocate memory\n");
 						goto error;
 					}
-					if(libusb_get_string_descriptor_ascii(out[found].handle, desc.iSerialNumber, (unsigned char*)serial, BUF_SIZE) < 0){
+					if(libusb_get_string_descriptor_ascii(handle, desc.iSerialNumber, (unsigned char*)&out->serial, BUF_SIZE) < 0){
 						fprintf(stderr, "LIBML: Cannot read device string descriptor\n");
 						goto error;
 					}
 #undef BUF_SIZE
-					out[found].serial = serial;
-					found++;
+                    return out;
 				}
 			}
 		}
-		out[found].handle = NULL;
-		out[found].serial = NULL;
 	}
 	libusb_free_device_list(device_list, 1);
-	return out;
+	return NULL;
 error:
-	if(res>0 && out){
-		for(ssize_t i=0; i<res; i++){
-			if(out[i].handle)
-				libusb_close(out[i].handle);
-			free(out[i].serial);
-		}
-	}
-	free(out);
+	if(out){
+        if(out->serial)
+            free(out->serial);
+        free(out);
+    }
+    if(handle)
+        libusb_close(handle);
 	libusb_free_device_list(device_list, 1);
-	return 0;
+	return NULL;
 }
 
 typedef struct{
@@ -140,7 +135,7 @@ int matelight_send_frame(matelight_handle *ml, void *buf, size_t w, size_t h, fl
 				return 1;
 			}
 			if(transferred != sizeof(frame)){
-				fprintf(stderr, "LIBML: Could not transfer all frame data. Only transferred %d out of %d bytes.\n", transferred, sizeof(frame));
+				fprintf(stderr, "LIBML: Could not transfer all frame data. Only transferred %d out of %zu bytes.\n", transferred, sizeof(frame));
 				return 1;
 			}
 		}
@@ -153,7 +148,7 @@ int matelight_send_frame(matelight_handle *ml, void *buf, size_t w, size_t h, fl
 		return 1;
 	}
 	if(transferred != sizeof(uint8_t)){
-		fprintf(stderr, "LIBML: Could not transfer all frame data. Only transferred %d out of %d bytes.\n", transferred, sizeof(uint8_t));
+		fprintf(stderr, "LIBML: Could not transfer all frame data. Only transferred %d out of %zu bytes.\n", transferred, sizeof(uint8_t));
 		return 1;
 	}
 	return 0;
diff --git a/host/usb.h b/host/usb.h
index fee3487..6e6aa8e 100644
--- a/host/usb.h
+++ b/host/usb.h
@@ -16,12 +16,12 @@
 
 typedef struct {
 	libusb_device_handle *handle;
-	char *serial;
+    char *serial;
 } matelight_handle;
 
 int matelight_usb_init(void);
 void matelight_usb_destroy(void);
-matelight_handle *matelight_open(void);
+matelight_handle *matelight_open(char *match_serial);
 int matelight_send_frame(matelight_handle *ml, void *buf, size_t w, size_t h, float brightness, int alpha);
 
 #endif//__USB_H__
diff --git a/host/viewer.py b/host/viewer.py
index 176c6d4..63a4546 100755
--- a/host/viewer.py
+++ b/host/viewer.py
@@ -24,7 +24,7 @@ if __name__ == '__main__':
 	udp_server = crap.CRAPServer(args.addr, args.port, blocking=True, log=lambda *_a: None)
 
 	with suppress(KeyboardInterrupt):
-		for frame in udp_server:
+		for _title, frame in udp_server:
 			bdf.printframe(frame)
 
 	udp_server.close()