new wedo plugin from TurtleBot

4 files changed, 264 insertions, 0 deletions

diff --git a/plugins/wedo_plugin/wedo/__init__.py b/plugins/wedo_plugin/wedo/__init__.py
new file mode 100644
index 0000000..8f77bb5
--- /dev/null
+++ b/plugins/wedo_plugin/wedo/__init__.py
@@ -0,0 +1,189 @@
+from functools import wraps
+
+from wedo.distance import interpolate_distance_data
+from wedo.motor import processMotorValues
+from wedo.tilt import process_tilt
+from wedo.tilt import FLAT, TILT_BACK, TILT_FORWARD, TILT_LEFT, TILT_RIGHT
+
+import usb.core
+import logging
+
+logger = logging.getLogger('wedo')
+
+ID_VENDOR, ID_PRODUCT = 0x0694, 0x0003
+UNAVAILABLE = None
+TILTSENSOR = (38, 39)
+DISTANCESENSOR = (176, 177, 178, 179)
+MOTOR = (0, 1, 2, 3, 238, 239)
+
+# limit the visibility to simplify the usage
+__all__ = ["scan_for_devices", "WeDo", "FLAT", "TILT_BACK", "TILT_FORWARD", "TILT_LEFT", "TILT_RIGHT"]
+
+def scan_for_devices():
+    """ Find all available devices """
+    return usb.core.find(find_all=True, idVendor=ID_VENDOR, idProduct=ID_PRODUCT)
+
+
+class WeDo(object):
+    """
+        Each instance of this class represents a physical WeDo device.
+
+        Usage :
+
+        >>> from wedo import WeDo
+        >>> wd = WeDo()
+
+        Activating the first motor full forward:
+        >>> wd.motor_a = 100
+
+        Activating the second motor half speed/force backward:
+        >>> wd.motor_b = -50
+
+        Current value of the tilt sensor:
+        >>> wd.tilt
+
+        Current distance value in meters of the distance sensor:
+        >>> wd.distance
+    """
+
+    def __init__(self, device=None):
+        """
+        If a device is not given, it will attach this instance to the first one found.
+        Otherwise you can pass a specific one from the list returned by scan_for_devices.
+        """
+        self.number = 0
+        self.dev = device
+        if self.dev is None:
+            devices = scan_for_devices()
+            if not devices:
+                raise OSError("Could not find a connected WeDo device")
+            self.dev = devices[0]
+        self.init_device()
+        self.valMotorA = 0
+        self.valMotorB = 0
+
+    def init_device(self):
+        """
+        Reinit device associated with the WeDo instance
+        """
+        if self.dev is None:
+            raise ValueError("No device attached to this instance")
+        try:
+            if self.dev.is_kernel_driver_active(0):
+                try:
+                    self.dev.detach_kernel_driver(0)
+                except usb.core.USBError as e:
+                    logger.error(
+                        "Could not detatch kernel driver: %s" % (str(e)))
+            self.endpoint = self.dev[0][(0, 0)][0]
+        except usb.core.USBError as e:
+            logger.error("Could not talk to WeDo device: %s" % (str(e)))
+
+    def getRawData(self):
+        """Read 64 bytes from the WeDo's endpoint, but only
+        return the last eight."""
+        try:
+            return self.endpoint.read(64)[-8:]
+        except usb.core.USBError as e:
+            logger.exception("Could not read from WeDo device")
+        return None
+
+    def setMotors(self):
+        """
+        Arguments should be in form of a number between 0
+        and 100, positive or negative. Magic numbers used for
+        the ctrl_transfer derived from sniffing USB coms.
+        """
+        data = [64, processMotorValues(self.valMotorA) & 0xFF, processMotorValues(self.valMotorB) & 0xFF,
+                0x00, 0x00, 0x00, 0x00, 0x00]
+        try:
+            self.dev.ctrl_transfer(bmRequestType=0x21, bRequest=0x09, wValue=0x0200, wIndex=0, data_or_wLength=data)
+        except usb.core.USBError as e:
+            logger.exception("Could not write to driver")
+
+    def getData(self):
+        """
+        Sensor data is contained in the 2nd and 4th byte, with
+        sensor IDs being contained in the 3rd and 5th byte
+        respectively.
+        """
+        rawData = self.getRawData()
+        if rawData is not None:
+            sensorData = {rawData[3]: rawData[2], rawData[5]: rawData[4]}
+        else:
+            sensorData = {}
+        return sensorData
+
+    @property
+    def raw_tilt(self):
+        """
+        Returns the raw tilt direction (arbitrary units)
+        """
+        data = self.getData()
+        for num in data:
+            if num in TILTSENSOR:
+                return data[num]
+        return UNAVAILABLE
+
+    @property
+    def tilt(self):
+        """
+        Returns the tilt direction (one of the FLAT, TILT_FORWARD, TILT_LEFT, TILT_RIGHT, TILT_BACK constants)
+        """
+        raw_data = self.raw_tilt
+        if raw_data is UNAVAILABLE:
+            return UNAVAILABLE
+        return process_tilt(raw_data)
+
+    @property
+    def raw_distance(self):
+        """
+        Return the raw evaluated distance from the distance meter (arbitrary units)
+        """
+        data = self.getData()
+        for num in data:
+            if num in DISTANCESENSOR:
+                return data[num]
+        return UNAVAILABLE
+
+    @property
+    def distance(self):
+        """
+        Return the evaluated distance in meters from the distance meter.
+        (Note: this is the ideal distance without any objets on the side, you might have to adapt it depending on your construction)
+        """
+
+        raw_data = self.raw_distance
+        if raw_data is UNAVAILABLE:
+            return UNAVAILABLE
+        return interpolate_distance_data(raw_data)
+
+    @property
+    def motor_a(self):
+        """ Get back the last speed/force set for motor A
+        """
+        return self.valMotorA
+
+    @property
+    def motor_b(self):
+        """ Get back the last speed/force set for motor A
+        """
+        return self.valMotorB
+
+    @motor_a.setter
+    def motor_a(self, value):
+        """ Sets the speed/force of the motor A, expects a value between -100 and 100
+        """
+        if value > 100 or value < -100:
+            raise ValueError("A motor can only be between -100 and 100")
+        self.valMotorA = value
+        self.setMotors()
+
+    @motor_b.setter
+    def motor_b(self, value):
+        """ Sets the speed/force of the motor B, expects a value between -100 and 100
+        """
+        if value > 100 or value < -100:
+            raise ValueError("A motor can only be between -100 and 100")
+        self.valMotorB = value
+        self.setMotors()
\ No newline at end of file
diff --git a/plugins/wedo_plugin/wedo/distance.py b/plugins/wedo_plugin/wedo/distance.py
new file mode 100644
index 0000000..0507b2c
--- /dev/null
+++ b/plugins/wedo_plugin/wedo/distance.py
@@ -0,0 +1,54 @@
+from bisect import bisect_left, bisect_right
+
+# data from distance sensor -> real measure in meters under ideal conditions (no side object, perpendicular wall)
+RAW_MEASURES = {
+    210: 0.46,
+    208: 0.39,
+    207: 0.34,
+    206: 0.32,
+    205: 0.305,
+    204: 0.29,
+    203: 0.27,
+    202: 0.26,
+    201: 0.25,
+    200: 0.245,
+    199: 0.235,
+    198: 0.225,
+    197: 0.22,
+    196: 0.215,
+    195: 0.2,
+    194: 0.195,
+    193: 0.18,
+    192: 0.175,
+    191: 0.17,
+    180: 0.15,
+    170: 0.13,
+    160: 0.125,
+    150: 0.11,
+    140: 0.105,
+    130: 0.1,
+    120: 0.095,
+    100: 0.075,
+    71: 0.065,
+    70: 0.06,
+    69: 0.053,
+    68: 0}
+
+RAW_MEASURES_KEYS = sorted(list(RAW_MEASURES))
+
+
+def interpolate_distance_data(raw):
+    left_index = bisect_left(RAW_MEASURES_KEYS, raw) - 1
+    if left_index < 0:
+        left_index = 0
+    right_index = left_index if left_index == len(RAW_MEASURES_KEYS) - 1 else left_index + 1
+
+    left = RAW_MEASURES_KEYS[left_index]
+    if left > raw:
+        return RAW_MEASURES[RAW_MEASURES_KEYS[left_index]]
+
+    right = RAW_MEASURES_KEYS[right_index]
+    mright = RAW_MEASURES[right]
+    mleft = RAW_MEASURES[left]
+    addup = ((raw - left) / (right - left)) * (mright - mleft) if mright != mleft else 0
+    return mleft + addup
\ No newline at end of file
diff --git a/plugins/wedo_plugin/wedo/motor.py b/plugins/wedo_plugin/wedo/motor.py
new file mode 100644
index 0000000..e74c00e
--- /dev/null
+++ b/plugins/wedo_plugin/wedo/motor.py
@@ -0,0 +1,7 @@
+def processMotorValues(value):
+    """Check to make sure motor values are sane."""
+    if 0 < value <= 100:
+        return value + 27
+    elif -100 <= value < 0:
+        return value - 27
+    return 0
diff --git a/plugins/wedo_plugin/wedo/tilt.py b/plugins/wedo_plugin/wedo/tilt.py
new file mode 100644
index 0000000..4f35340
--- /dev/null
+++ b/plugins/wedo_plugin/wedo/tilt.py
@@ -0,0 +1,14 @@
+FLAT, TILT_FORWARD, TILT_LEFT, TILT_RIGHT, TILT_BACK = range(5)
+
+def process_tilt(raw_tilt_value):
+    """Use a series of elif/value-checks to process the tilt
+    sensor data."""
+    if 10 <= raw_tilt_value <= 40:
+        return TILT_BACK
+    elif 60 <= raw_tilt_value <= 90:
+        return TILT_RIGHT
+    elif 170 <= raw_tilt_value <= 190:
+        return TILT_FORWARD
+    elif 220 <= raw_tilt_value <= 240:
+        return TILT_LEFT
+    return FLAT