pybricks/util_mp: Utility for percentage array.

This will be useful in several places, so make a utility function.

Now it also accepts a single floating point value for consistency with the rest of the API.

Fixes pybricks/support#1547.

Author: laurensvalk

CHANGELOG.md

@@ -4,6 +4,11 @@
 
 ## [Unreleased]
 
+### Changed
+- Allow single floating point value for brightness array ([support#1547]).
+
+[support#1547]: https://github.com/pybricks/support/issues/1547
+
 ## [3.4.0] - 2024-03-11
 
 ### Changed

pybricks/common/pb_type_lightarray.c

@@ -32,26 +32,7 @@ STATIC mp_obj_t common_LightArray_on(size_t n_args, const mp_obj_t *pos_args, mp
         PB_ARG_DEFAULT_INT(brightness, 100));
 
     int8_t brightness_values[4];
-
-    // Given an integer, make all lights the same brightness.
-    if (mp_obj_is_int(brightness_in)) {
-        int32_t b = pb_obj_get_pct(brightness_in);
-        for (uint8_t i = 0; i < self->number_of_lights; i++) {
-            brightness_values[i] = b;
-        }
-    }
-    // Otherwise, get each brightness value from list or tuple.
-    else {
-        mp_obj_t *brightness_objects;
-        size_t num_values;
-        mp_obj_get_array(brightness_in, &num_values, &brightness_objects);
-        if (num_values != self->number_of_lights) {
-            pb_assert(PBIO_ERROR_INVALID_ARG);
-        }
-        for (uint8_t i = 0; i < self->number_of_lights; i++) {
-            brightness_values[i] = pb_obj_get_pct(brightness_objects[i]);
-        }
-    }
+    pb_obj_get_pct_or_array(brightness_in, self->number_of_lights, brightness_values);
 
     // Set the brightness values and wait or await it.
     return pb_type_device_set_data(self->sensor, self->light_mode, brightness_values, self->number_of_lights);

pybricks/util_mp/pb_obj_helper.c

@@ -124,6 +124,34 @@ mp_int_t pb_obj_get_default_abs_int(mp_obj_t obj, mp_int_t default_val) {
     return value > 0 ? value: -value;
 }
 
+/**
+ * Populates an array of percentages from a single number or tuple.
+ *
+ * @param obj_in [in]  A MicroPython object
+ * @param num    [in]  The number of values to populate
+ * @param values [out] The array to populate
+ *
+ * Raises exception if @p obj_in is not a number or a tuple of @p num values.
+ */
+void pb_obj_get_pct_or_array(mp_obj_t obj_in, size_t num, int8_t *values) {
+    if (mp_obj_is_type(obj_in, &mp_type_tuple) || mp_obj_is_type(obj_in, &mp_type_list)) {
+        mp_obj_t *items;
+        size_t len;
+        mp_obj_get_array(obj_in, &len, &items);
+        if (len != num) {
+            mp_raise_msg_varg(&mp_type_ValueError, MP_ERROR_TEXT("expected 1 number or tuple of %d values"), num);
+        }
+        for (size_t i = 0; i < num; i++) {
+            values[i] = pb_obj_get_pct(items[i]);
+        }
+    } else {
+        mp_int_t pct = pb_obj_get_pct(obj_in);
+        for (size_t i = 0; i < num; i++) {
+            values[i] = pct;
+        }
+    }
+}
+
 mp_obj_t pb_obj_get_base_class_obj(mp_obj_t obj, const mp_obj_type_t *type) {
 
     // If it equals the base type then return as is

pybricks/util_mp/pb_obj_helper.h

@@ -30,6 +30,9 @@ mp_obj_t pb_obj_new_fraction(int32_t numerator, int32_t denominator);
 // Get absolute value if object is not none, else return default
 mp_int_t pb_obj_get_default_abs_int(mp_obj_t obj, mp_int_t default_val);
 
+// Get array of percentages from single value or tuple
+void pb_obj_get_pct_or_array(mp_obj_t obj_in, size_t num, int8_t *values);
+
 // Get base instance if object is instance of subclass of type
 mp_obj_t pb_obj_get_base_class_obj(mp_obj_t obj, const mp_obj_type_t *type);