lern-tracking-system/raspberry-pi/connection.py

import threading, time, serial

class SerialConnection:
  def __init__(self):
    self._ser = None
    self.port = None
    
  def open(self, port = None, baudrate = 1000000):
    self.port = port

    if self.port == None:
      for port in ["/dev/ttyUSB{}".format(p) for p in range(4)]:
        try:
          self._ser = serial.Serial(port, baudrate)
          print("connected to " + port)
          self.port = port
          break
        except serial.SerialException:
          pass
      
    if self.port == None:
      for port in ["COM{}".format(p) for p in range(3,20)]:
        try:
          self._ser = serial.Serial(port, baudrate)
          print("connected to " + port)
          self.port = port
          break
        except serial.SerialException:
          pass
    else:
      try:
        self._ser = serial.Serial(self.port, baudrate)
        print("connected to " + self.port)
      except serial.SerialException:
        pass

    if not self._ser:
        print("connection failed")
    
  def send(self, bytes):
    if self._ser:
      try:
        return self._ser.write(bytes)
      except serial.SerialException:
        self._ser = None
    if not self._ser:
      self.open()
      if not self._ser:
        time.sleep(30)
          
  def read(self, n=64):
    if self._ser:
      try:
        if self._ser.in_waiting:
          return self._ser.read(n)
      except serial.SerialException:
        self._ser = None
    if not self._ser:
      self.open()

  def readline(self):
    if self._ser:
      return self._ser.readline()
    else:
      return None

  def isConnected(self):
    return self._ser != None

  def close(self):
    if self._ser:
      self._ser.close()
      self._ser = None

class ArduinoSlave(SerialConnection):
  def __init__(self):
    super().__init__()
    self.sensorData = [0] * 13
    self._recvCbs = []
    self._t = None

  def open(self, port = None, baudrate = 1000000):
    super().open(port, baudrate)
    if not self._t:
      self._t = threading.Thread(target=self._readSensors, args=())
      self._t.daemon = True  # thread dies when main thread (only non-daemon thread) exits.
      self._t.start()

  def close(self):
    super().close()
    if self._t:
      self._t._stop()

  def getAcusticRTTs(self): # in microseconds
    return [
      int(self.sensorData[0]) / 16,
      int(self.sensorData[1]) / 16
    ]

  def getMagneticField(self): # in mT
    return [
      int(self.sensorData[2]) / 1000,
      int(self.sensorData[3]) / 1000,
      int(self.sensorData[4]) / 1000
    ]
  
  def getAccelValues(self):
    return [
      int(self.sensorData[5]) / 1000,
      int(self.sensorData[6]) / 1000,
      int(self.sensorData[7]) / 1000
    ]
  
  def getGyroValues(self):
    return [
      int(self.sensorData[8])  / 1000,
      int(self.sensorData[9])  / 1000,
      int(self.sensorData[10]) / 1000
    ]
  
  def getTemperature(self): # in °C
    return int(self.sensorData[11]) / 1000

  def addRecvCallback(self, cb):
    self._recvCbs.append(cb)

  def _readSensors(self):
    while True:
      data = super().readline()
      if data:
        data = str(data, "ASCII")
        vals = data.split('\t')
        if vals[0] == "DATA:" and len(vals) >= 13:
          self.sensorData = vals[1:]
          for cb in self._recvCbs:
            cb(self.sensorData)
        else:
          print("SERIAL: ", data[:-2])

if __name__ == "__main__":
  arduino = ArduinoSlave()
  def cb(x):
    print(arduino.getAcusticRTTs(), arduino.getMagneticField(), arduino.getTemperature())

  arduino.addRecvCallback(cb)
  arduino.open()
  while True:
    time.sleep(1)