Add indices() support to slice.

1 files changed, 83 insertions, 0 deletions

diff --git a/py/objslice.c b/py/objslice.c
index 928be6dab..6c9592628 100644
--- a/py/objslice.c
+++ b/py/objslice.c
@@ -29,6 +29,7 @@
 
 #include "py/nlr.h"
 #include "py/obj.h"
+#include "py/runtime.h"
 #include "py/runtime0.h"
 
 /******************************************************************************/
@@ -58,6 +59,58 @@ STATIC void slice_print(const mp_print_t *print, mp_obj_t o_in, mp_print_kind_t
 }
 
 #if MICROPY_PY_BUILTINS_SLICE_ATTRS
+STATIC mp_obj_t slice_indices(mp_obj_t self_in, mp_obj_t length_obj) {
+    mp_obj_slice_t *self = MP_OBJ_TO_PTR(self_in);
+    if (!MP_OBJ_IS_SMALL_INT(length_obj)) {
+        mp_raise_TypeError("Length must be an int");
+    }
+
+    int length = MP_OBJ_SMALL_INT_VALUE(length_obj);
+    if (length < 0) {
+        mp_raise_ValueError("Length must be non-negative");
+    }
+
+    mp_obj_t indices[3] = {MP_OBJ_NEW_SMALL_INT(0), length_obj, MP_OBJ_NEW_SMALL_INT(1)};
+    mp_obj_t slice[2] = {self->start, self->stop};
+
+    int step = 1;
+    if (self->step != mp_const_none) {
+        indices[2] = self->step;
+        step = MP_OBJ_SMALL_INT_VALUE(self->step);
+        if (step < 0) {
+            indices[0] = MP_OBJ_NEW_SMALL_INT(length - 1);
+            indices[1] = MP_OBJ_NEW_SMALL_INT(-1);
+        }
+        if (step == 0) {
+            mp_raise_ValueError("slice step cannot be zero");
+        }
+    }
+    for (int i = 0; i < 2; i++) {
+        if (slice[i] == mp_const_none) {
+            continue;
+        }
+        int value = MP_OBJ_SMALL_INT_VALUE(slice[i]);
+        if (value < 0) {
+            value += length;
+        }
+        if (value < 0) {
+            if (step > 0) {
+                value = 0;
+            } else if (step < 0) {
+                value = -1;
+            }
+        } else if (value > length) {
+            value = length;
+        }
+        indices[i] = MP_OBJ_NEW_SMALL_INT(value);
+    }
+
+    mp_obj_t tuple = mp_obj_new_tuple(3, indices);
+
+    return tuple;
+}
+MP_DEFINE_CONST_FUN_OBJ_2(slice_indices_obj, slice_indices);
+
 STATIC void slice_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
     if (dest[0] != MP_OBJ_NULL) {
         // not load attribute
@@ -70,8 +123,13 @@ STATIC void slice_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
         dest[0] = self->stop;
     } else if (attr == MP_QSTR_step) {
         dest[0] = self->step;
+    } else if (attr == MP_QSTR_indices) {
+        mp_convert_member_lookup(self_in, self->base.type, (mp_obj_t) &slice_indices_obj, dest);
     }
 }
+
+STATIC mp_obj_t slice_make_new(const mp_obj_type_t *type,
+        mp_uint_t n_args, mp_uint_t n_kw, const mp_obj_t *args);
 #endif
 
 const mp_obj_type_t mp_type_slice = {
@@ -79,6 +137,7 @@ const mp_obj_type_t mp_type_slice = {
     .name = MP_QSTR_slice,
     .print = slice_print,
 #if MICROPY_PY_BUILTINS_SLICE_ATTRS
+    .make_new = slice_make_new,
     .attr = slice_attr,
 #endif
 };
@@ -92,6 +151,30 @@ mp_obj_t mp_obj_new_slice(mp_obj_t ostart, mp_obj_t ostop, mp_obj_t ostep) {
     return MP_OBJ_FROM_PTR(o);
 }
 
+#if MICROPY_PY_BUILTINS_SLICE_ATTRS
+STATIC mp_obj_t slice_make_new(const mp_obj_type_t *type,
+        mp_uint_t n_args, mp_uint_t n_kw, const mp_obj_t *args) {
+    // check number of arguments
+    mp_arg_check_num(n_args, n_kw, 1, 3, false);
+
+    // 1st argument is the pin
+    mp_obj_t start = mp_const_none;
+    mp_obj_t stop = mp_const_none;
+    mp_obj_t step = mp_const_none;
+    if (n_args == 1) {
+        stop = args[0];
+    } else {
+        start = args[0];
+        stop = args[1];
+        if (n_args == 3) {
+            step = args[2];
+        }
+    }
+
+    return mp_obj_new_slice(start, stop, step);
+}
+#endif
+
 void mp_obj_slice_get(mp_obj_t self_in, mp_obj_t *start, mp_obj_t *stop, mp_obj_t *step) {
     assert(MP_OBJ_IS_TYPE(self_in, &mp_type_slice));
     mp_obj_slice_t *self = MP_OBJ_TO_PTR(self_in);