extmod/modpupdevices: Add InfraredSensor

Using the raw data, we can get much more information than the 10 distance increments that this sensor normally gives.

We provide two convenience methods based on this data:
- distance() nonlinearly scales the raw data to provide an approximately
  linear distance value. It corresponds to the original sensor value,
  but with higher resolution and a scale from 0-100 to match similar
  relative distance methods on the other sensors. Resolution is somewhat
  course above 50%.
- reflection() linearly scales the data to provide reflection data close
  to the sensor. This can be used for line following or precise distance
  sensing up close.

There is also a count() method, which does the same as the original
sensor. In this class, the count is reset to 0 upon instantiation. Using
the other methods does not reset this count.

Author: laurensvalk

extmod/modpupdevices.c

@@ -526,6 +526,81 @@ STATIC const mp_obj_type_t pupdevices_ForceSensor_type = {
     .locals_dict = (mp_obj_dict_t *)&pupdevices_ForceSensor_locals_dict,
 };
 
+// Class structure for InfraredSensor
+typedef struct _pupdevices_InfraredSensor_obj_t {
+    mp_obj_base_t base;
+    pbdevice_t *pbdev;
+    int32_t count_offset;
+} pupdevices_InfraredSensor_obj_t;
+
+// pybricks.pupdevices.InfraredSensor._raw
+STATIC int32_t pupdevices_InfraredSensor__raw(pbdevice_t *pbdev) {
+    int32_t raw[3];
+    pbdevice_get_values(pbdev, PBIO_IODEV_MODE_PUP_WEDO2_MOTION_SENSOR__CAL, raw);
+    return raw[0];
+}
+
+// pybricks.pupdevices.InfraredSensor.__init__
+STATIC mp_obj_t pupdevices_InfraredSensor_make_new(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args) {
+    PB_PARSE_ARGS_CLASS(n_args, n_kw, args,
+        PB_ARG_REQUIRED(port));
+
+    pupdevices_InfraredSensor_obj_t *self = m_new_obj(pupdevices_InfraredSensor_obj_t);
+    self->base.type = (mp_obj_type_t *)type;
+
+    mp_int_t port_num = pb_type_enum_get_value(port, &pb_enum_type_Port);
+
+    // Get iodevice
+    self->pbdev = pbdevice_get_device(port_num, PBIO_IODEV_TYPE_ID_WEDO2_MOTION_SENSOR);
+
+    // Reset sensor counter and get sensor back in sensing mode
+    pbdevice_get_values(self->pbdev, PBIO_IODEV_MODE_PUP_WEDO2_MOTION_SENSOR__COUNT, &self->count_offset);
+    pupdevices_InfraredSensor__raw(self->pbdev);
+
+    return MP_OBJ_FROM_PTR(self);
+}
+
+// pybricks.pupdevices.InfraredSensor.count
+STATIC mp_obj_t pupdevices_InfraredSensor_count(mp_obj_t self_in) {
+    pupdevices_InfraredSensor_obj_t *self = MP_OBJ_TO_PTR(self_in);
+    int32_t count;
+    pbdevice_get_values(self->pbdev, PBIO_IODEV_MODE_PUP_WEDO2_MOTION_SENSOR__COUNT, &count);
+    return mp_obj_new_int(count - self->count_offset);
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_1(pupdevices_InfraredSensor_count_obj, pupdevices_InfraredSensor_count);
+
+// pybricks.pupdevices.InfraredSensor.reflection
+STATIC mp_obj_t pupdevices_InfraredSensor_reflection(mp_obj_t self_in) {
+    pupdevices_InfraredSensor_obj_t *self = MP_OBJ_TO_PTR(self_in);
+    int32_t raw = pupdevices_InfraredSensor__raw(self->pbdev);
+    return mp_obj_new_int(raw / 5);
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_1(pupdevices_InfraredSensor_reflection_obj, pupdevices_InfraredSensor_reflection);
+
+// pybricks.pupdevices.InfraredSensor.distance
+STATIC mp_obj_t pupdevices_InfraredSensor_distance(mp_obj_t self_in) {
+    pupdevices_InfraredSensor_obj_t *self = MP_OBJ_TO_PTR(self_in);
+    int32_t raw = pupdevices_InfraredSensor__raw(self->pbdev);
+    return mp_obj_new_int(1100 / (10 + raw));
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_1(pupdevices_InfraredSensor_distance_obj, pupdevices_InfraredSensor_distance);
+
+// dir(pybricks.pupdevices.InfraredSensor)
+STATIC const mp_rom_map_elem_t pupdevices_InfraredSensor_locals_dict_table[] = {
+    { MP_ROM_QSTR(MP_QSTR_count),       MP_ROM_PTR(&pupdevices_InfraredSensor_count_obj)                },
+    { MP_ROM_QSTR(MP_QSTR_reflection),  MP_ROM_PTR(&pupdevices_InfraredSensor_reflection_obj)           },
+    { MP_ROM_QSTR(MP_QSTR_distance),    MP_ROM_PTR(&pupdevices_InfraredSensor_distance_obj)             },
+};
+STATIC MP_DEFINE_CONST_DICT(pupdevices_InfraredSensor_locals_dict, pupdevices_InfraredSensor_locals_dict_table);
+
+// type(pybricks.pupdevices.InfraredSensor)
+STATIC const mp_obj_type_t pupdevices_InfraredSensor_type = {
+    { &mp_type_type },
+    .name = MP_QSTR_InfraredSensor,
+    .make_new = pupdevices_InfraredSensor_make_new,
+    .locals_dict = (mp_obj_dict_t *)&pupdevices_InfraredSensor_locals_dict,
+};
+
 // Class structure for Light
 typedef struct _pupdevices_Light_obj_t {
     mp_obj_base_t base;
@@ -593,6 +668,7 @@ STATIC const mp_rom_map_elem_t pupdevices_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR_ColorSensor),         MP_ROM_PTR(&pupdevices_ColorSensor_type)         },
     { MP_ROM_QSTR(MP_QSTR_UltrasonicSensor),    MP_ROM_PTR(&pupdevices_UltrasonicSensor_type)    },
     { MP_ROM_QSTR(MP_QSTR_ForceSensor),         MP_ROM_PTR(&pupdevices_ForceSensor_type)         },
+    { MP_ROM_QSTR(MP_QSTR_InfraredSensor),      MP_ROM_PTR(&pupdevices_InfraredSensor_type)         },
     { MP_ROM_QSTR(MP_QSTR_Light),               MP_ROM_PTR(&pupdevices_Light_type)               },
 };
 

extmod/pbdevice.h

@@ -118,6 +118,13 @@ enum {
     PBIO_IODEV_MODE_NXT_ENERGY_METER_ALL           = 7,
 };
 
+// LEGO POWERED UP WEDO 2.0 Infrared "Motion" sensor
+enum {
+    PBIO_IODEV_MODE_PUP_WEDO2_MOTION_SENSOR__DETECT = 0,  // read 1x int8_t
+    PBIO_IODEV_MODE_PUP_WEDO2_MOTION_SENSOR__COUNT  = 1,  // read 1x int32_t
+    PBIO_IODEV_MODE_PUP_WEDO2_MOTION_SENSOR__CAL    = 2,  // read 3x int16_t
+};
+
 // LEGO POWERED UP Color and Distance Sensor
 enum {
     PBIO_IODEV_MODE_PUP_COLOR_DISTANCE_SENSOR__COLOR = 0,  // read 1x int8_t