Skip to content

Commit bb91f11

Browse files
committed
py: Improve range: add len, subscr, proper print.
Can now index ranges with integers and slices, and reverse ranges (although reversing is not very efficient). Not sure how useful this stuff is, but gets us closer to having all of Python's builtins.
1 parent 4c03b3a commit bb91f11

3 files changed

Lines changed: 92 additions & 2 deletions

File tree

py/objrange.c

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
#include "misc.h"
3232
#include "qstr.h"
3333
#include "obj.h"
34+
#include "runtime0.h"
3435
#include "runtime.h"
3536

3637
/******************************************************************************/
@@ -82,6 +83,16 @@ typedef struct _mp_obj_range_t {
8283
mp_int_t step;
8384
} mp_obj_range_t;
8485

86+
STATIC void range_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t self_in, mp_print_kind_t kind) {
87+
mp_obj_range_t *self = self_in;
88+
print(env, "range(%d, %d", self->start, self->stop);
89+
if (self->step == 1) {
90+
print(env, ")");
91+
} else {
92+
print(env, ", %d)", self->step);
93+
}
94+
}
95+
8596
STATIC mp_obj_t range_make_new(mp_obj_t type_in, uint n_args, uint n_kw, const mp_obj_t *args) {
8697
mp_arg_check_num(n_args, n_kw, 1, 3, false);
8798

@@ -96,13 +107,63 @@ STATIC mp_obj_t range_make_new(mp_obj_t type_in, uint n_args, uint n_kw, const m
96107
o->start = mp_obj_get_int(args[0]);
97108
o->stop = mp_obj_get_int(args[1]);
98109
if (n_args == 3) {
110+
// TODO check step is non-zero
99111
o->step = mp_obj_get_int(args[2]);
100112
}
101113
}
102114

103115
return o;
104116
}
105117

118+
STATIC mp_int_t range_len(mp_obj_range_t *self) {
119+
// When computing length, need to take into account step!=1 and step<0.
120+
mp_int_t len = self->stop - self->start + self->step;
121+
if (self->step > 0) {
122+
len -= 1;
123+
} else {
124+
len += 1;
125+
}
126+
len = len / self->step;
127+
if (len < 0) {
128+
len = 0;
129+
}
130+
return len;
131+
}
132+
133+
STATIC mp_obj_t range_unary_op(int op, mp_obj_t self_in) {
134+
mp_obj_range_t *self = self_in;
135+
mp_int_t len = range_len(self);
136+
switch (op) {
137+
case MP_UNARY_OP_BOOL: return MP_BOOL(len > 0);
138+
case MP_UNARY_OP_LEN: return MP_OBJ_NEW_SMALL_INT(len);
139+
default: return MP_OBJ_NULL; // op not supported
140+
}
141+
}
142+
143+
STATIC mp_obj_t range_subscr(mp_obj_t self_in, mp_obj_t index, mp_obj_t value) {
144+
if (value == MP_OBJ_SENTINEL) {
145+
// load
146+
mp_obj_range_t *self = self_in;
147+
mp_int_t len = range_len(self);
148+
#if MICROPY_PY_BUILTINS_SLICE
149+
if (MP_OBJ_IS_TYPE(index, &mp_type_slice)) {
150+
mp_bound_slice_t slice;
151+
mp_seq_get_fast_slice_indexes(len, index, &slice);
152+
mp_obj_range_t *o = m_new_obj(mp_obj_range_t);
153+
o->base.type = &mp_type_range;
154+
o->start = slice.start;
155+
o->stop = slice.stop;
156+
o->step = slice.step;
157+
return o;
158+
}
159+
#endif
160+
uint index_val = mp_get_index(self->base.type, len, index, false);
161+
return MP_OBJ_NEW_SMALL_INT(self->start + index_val * self->step);
162+
} else {
163+
return MP_OBJ_NULL; // op not supported
164+
}
165+
}
166+
106167
STATIC mp_obj_t range_getiter(mp_obj_t o_in) {
107168
mp_obj_range_t *o = o_in;
108169
return mp_obj_new_range_iterator(o->start, o->stop, o->step);
@@ -111,6 +172,9 @@ STATIC mp_obj_t range_getiter(mp_obj_t o_in) {
111172
const mp_obj_type_t mp_type_range = {
112173
{ &mp_type_type },
113174
.name = MP_QSTR_range,
175+
.print = range_print,
114176
.make_new = range_make_new,
177+
.unary_op = range_unary_op,
178+
.subscr = range_subscr,
115179
.getiter = range_getiter,
116180
};

tests/basics/builtin_range.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
# test builtin range type
2+
3+
# print
4+
print(range(4))
5+
6+
# len
7+
print(len(range(0)))
8+
print(len(range(4)))
9+
print(len(range(1, 4)))
10+
print(len(range(1, 4, 2)))
11+
print(len(range(1, 4, -1)))
12+
print(len(range(4, 1, -1)))
13+
print(len(range(4, 1, -2)))
14+
15+
# subscr
16+
print(range(4)[0])
17+
print(range(4)[1])
18+
print(range(4)[-1])
19+
20+
# slice
21+
print(range(4)[0:])
22+
print(range(4)[1:])
23+
print(range(4)[1:2])
24+
print(range(4)[1:3])
25+
print(range(4)[1::2])
26+
print(range(4)[1:-2:2])

tests/basics/builtin_reversed.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,8 @@
1818
print(b)
1919

2020
# range
21-
#for i in reversed(range(3)):
22-
# print(i)
21+
for i in reversed(range(3)):
22+
print(i)
2323

2424
# user object
2525
class A:

0 commit comments

Comments
 (0)