extmod/modure: Add match.span(), start() and end() methods, and tests.

This feature is controlled at compile time by
MICROPY_PY_URE_MATCH_SPAN_START_END, disabled by default.

Thanks to @dmazzella for the original patch for this feature; see #3770.

1 files changed, 55 insertions, 0 deletions

diff --git a/extmod/modure.c b/extmod/modure.c
index aec7350b2..a536f907f 100644
--- a/extmod/modure.c
+++ b/extmod/modure.c
@@ -94,11 +94,66 @@ MP_DEFINE_CONST_FUN_OBJ_1(match_groups_obj, match_groups);
 
 #endif
 
+#if MICROPY_PY_URE_MATCH_SPAN_START_END
+
+STATIC void match_span_helper(size_t n_args, const mp_obj_t *args, mp_obj_t span[2]) {
+    mp_obj_match_t *self = MP_OBJ_TO_PTR(args[0]);
+
+    mp_int_t no = 0;
+    if (n_args == 2) {
+        no = mp_obj_get_int(args[1]);
+        if (no < 0 || no >= self->num_matches) {
+            nlr_raise(mp_obj_new_exception_arg1(&mp_type_IndexError, args[1]));
+        }
+    }
+
+    mp_int_t s = -1;
+    mp_int_t e = -1;
+    const char *start = self->caps[no * 2];
+    if (start != NULL) {
+        // have a match for this group
+        const char *begin = mp_obj_str_get_str(self->str);
+        s = start - begin;
+        e = self->caps[no * 2 + 1] - begin;
+    }
+
+    span[0] = mp_obj_new_int(s);
+    span[1] = mp_obj_new_int(e);
+}
+
+STATIC mp_obj_t match_span(size_t n_args, const mp_obj_t *args) {
+    mp_obj_t span[2];
+    match_span_helper(n_args, args, span);
+    return mp_obj_new_tuple(2, span);
+}
+MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(match_span_obj, 1, 2, match_span);
+
+STATIC mp_obj_t match_start(size_t n_args, const mp_obj_t *args) {
+    mp_obj_t span[2];
+    match_span_helper(n_args, args, span);
+    return span[0];
+}
+MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(match_start_obj, 1, 2, match_start);
+
+STATIC mp_obj_t match_end(size_t n_args, const mp_obj_t *args) {
+    mp_obj_t span[2];
+    match_span_helper(n_args, args, span);
+    return span[1];
+}
+MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(match_end_obj, 1, 2, match_end);
+
+#endif
+
 STATIC const mp_rom_map_elem_t match_locals_dict_table[] = {
     { MP_ROM_QSTR(MP_QSTR_group), MP_ROM_PTR(&match_group_obj) },
     #if MICROPY_PY_URE_MATCH_GROUPS
     { MP_ROM_QSTR(MP_QSTR_groups), MP_ROM_PTR(&match_groups_obj) },
     #endif
+    #if MICROPY_PY_URE_MATCH_SPAN_START_END
+    { MP_ROM_QSTR(MP_QSTR_span), MP_ROM_PTR(&match_span_obj) },
+    { MP_ROM_QSTR(MP_QSTR_start), MP_ROM_PTR(&match_start_obj) },
+    { MP_ROM_QSTR(MP_QSTR_end), MP_ROM_PTR(&match_end_obj) },
+    #endif
 };
 
 STATIC MP_DEFINE_CONST_DICT(match_locals_dict, match_locals_dict_table);