if __name__ == '__main__': import sys sys.exit('This is a library file to be imported into your own python scripts. It doesn\'t do anything if run directly') from cobs import cobs from collections import ( defaultdict, ) from contextlib import contextmanager from enum import Enum import logging import os from queue import ( Empty, Full, Queue, ) from random import randint import serial import serial.tools.list_ports import sys from threading import ( Thread, Lock, ) import time import zlib software_root = os.path.dirname(os.path.abspath(__file__)) sys.path.append(os.path.join(software_root, 'proto_gen')) from proto_gen import smartknob_pb2 SMARTKNOB_BAUD = 921600 class Smartknob(object): RETRY_TIMEOUT = 0.25 def __init__(self, serial_instance): self._serial = serial_instance self._logger = logging.getLogger('smartknob') self._out_q = Queue() self._ack_q = Queue() self._next_nonce = randint(0, 255) self._run = True self._lock = Lock() self._message_handlers = defaultdict(list) def _read_loop(self): self._logger.debug('Read loop started') buffer = b'' while True: buffer += self._serial.read_until(b'\0') if not self._run: return if not len(buffer): continue if not buffer.endswith(b'\0'): continue self._process_frame(buffer[:-1]) buffer = b'' def _process_frame(self, frame): try: decoded = cobs.decode(frame) except cobs.DecodeError: self._logger.debug(f'Failed decode ({len(frame)} bytes)') self._logger.debug(frame) return if len(decoded) < 4: return payload = decoded[:-4] expected_crc = zlib.crc32(payload) & 0xffffffff provided_crc = (decoded[-1] << 24) \ | (decoded[-2] << 16) \ | (decoded[-3] << 8) \ | decoded[-4] if expected_crc != provided_crc: self._logger.debug(f'Bad CRC. expected={hex(expected_crc)}, actual={hex(provided_crc)}') return message = smartknob_pb2.FromSmartKnob() message.ParseFromString(payload) self._logger.debug(message) payload_type = message.WhichOneof('payload') # If this is an ack, notify the write thread if payload_type == 'ack': nonce = message.ack.nonce self._ack_q.put(nonce) with self._lock: for handler in self._message_handlers[payload_type] + self._message_handlers[None]: try: handler(getattr(message, payload_type)) except: self._logger.warning(f'Unhandled exception in message handler ({payload_type})', exc_info=True) def _write_loop(self): self._logger.debug('Write loop started') while True: data = self._out_q.get() # Check for shutdown if not self._run: self._logger.debug('Write loop exiting @ _out_q') return (nonce, encoded_message) = data next_retry = 0 while True: if time.time() >= next_retry: if next_retry > 0: self._logger.debug('Retry write...') self._serial.write(encoded_message) self._serial.write(b'\0') next_retry = time.time() + Smartknob.RETRY_TIMEOUT try: latest_ack_nonce = self._ack_q.get(timeout=next_retry - time.time()) except Empty: latest_ack_nonce = None # Check for shutdown if not self._run: self._logger.debug('Write loop exiting @ _ack_q') return if latest_ack_nonce == nonce: break else: self._logger.debug(f'Got unexpected nonce: {latest_ack_nonce}') def _enqueue_message(self, message): nonce = self._next_nonce self._next_nonce += 1 message.nonce = nonce payload = bytearray(message.SerializeToString()) crc = zlib.crc32(payload) & 0xffffffff payload.append(crc & 0xff) payload.append((crc >> 8) & 0xff) payload.append((crc >> 16) & 0xff) payload.append((crc >> 24) & 0xff) encoded_message = cobs.encode(payload) self._out_q.put((nonce, encoded_message)) approx_q_length = self._out_q.qsize() self._logger.debug(f'Out q length: {approx_q_length}') if approx_q_length > 10: self._logger.warning(f'Output queue length is high! ({approx_q_length}) Is the smartknob still connected and functional?') def set_config(self, config): message = smartknob_pb2.ToSmartknob() message.smartknob_config.CopyFrom(config) self._enqueue_message(message) def start(self): self.read_thread = Thread(target=self._read_loop) self.write_thread = Thread(target=self._write_loop) self.read_thread.start() self.write_thread.start() def shutdown(self): self._logger.info('Shutting down...') self._run = False self.read_thread.join() self._logger.debug('Read thread terminated') self._out_q.put(None) self._ack_q.put(None) self.write_thread.join() self._logger.debug('Write thread terminated') def add_handler(self, message_type, handler): with self._lock: self._message_handlers[message_type].append(handler) return lambda: self._remove_handler(message_type, handler) def _remove_handler(self, message_type, handler): with self._lock: self._message_handlers[message_type].remove(handler) def request_state(self): message = smartknob_pb2.ToSmartknob() message.request_state.SetInParent() self._enqueue_message(message) def hard_reset(self): self._serial.setRTS(True) self._serial.setDTR(False) time.sleep(0.2) self._serial.setDTR(True) time.sleep(0.2) @contextmanager def smartknob_context(serial_port, default_logging=True, wait_for_comms=True): with serial.Serial(serial_port, SMARTKNOB_BAUD, timeout=1.0) as ser: s = Smartknob(ser) s.start() if default_logging: s.add_handler('log', lambda msg: s._logger.info(f'From smartknob: {msg.msg}')) if wait_for_comms: s._logger.info('Connecting to smartknob...') q = Queue(1) def startup_handler(message): try: q.put_nowait(None) except Full: pass unregister = s.add_handler('smartknob_state', startup_handler) s.request_state() q.get() unregister() s._logger.info('Connected!') try: yield s finally: s.shutdown() def ask_for_serial_port(): """ Helper function to ask which port to use via stdin """ print('Available ports:') ports = sorted( filter( lambda p: p.description != 'n/a', serial.tools.list_ports.comports(), ), key=lambda p: p.device, ) for i, port in enumerate(ports): print('[{: 2}] {} - {}'.format(i, port.device, port.description)) print() value = input('Use which port? ') port_index = int(value) assert 0 <= port_index < len(ports) return ports[port_index].device