Commit bb91f119 authored by Damien George's avatar Damien George
Browse files

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.
parent 4c03b3a8
......@@ -31,6 +31,7 @@
#include "misc.h"
#include "qstr.h"
#include "obj.h"
#include "runtime0.h"
#include "runtime.h"
/******************************************************************************/
......@@ -82,6 +83,16 @@ typedef struct _mp_obj_range_t {
mp_int_t step;
} mp_obj_range_t;
STATIC void range_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t self_in, mp_print_kind_t kind) {
mp_obj_range_t *self = self_in;
print(env, "range(%d, %d", self->start, self->stop);
if (self->step == 1) {
print(env, ")");
} else {
print(env, ", %d)", self->step);
}
}
STATIC mp_obj_t range_make_new(mp_obj_t type_in, uint n_args, uint n_kw, const mp_obj_t *args) {
mp_arg_check_num(n_args, n_kw, 1, 3, false);
......@@ -96,6 +107,7 @@ STATIC mp_obj_t range_make_new(mp_obj_t type_in, uint n_args, uint n_kw, const m
o->start = mp_obj_get_int(args[0]);
o->stop = mp_obj_get_int(args[1]);
if (n_args == 3) {
// TODO check step is non-zero
o->step = mp_obj_get_int(args[2]);
}
}
......@@ -103,6 +115,55 @@ STATIC mp_obj_t range_make_new(mp_obj_t type_in, uint n_args, uint n_kw, const m
return o;
}
STATIC mp_int_t range_len(mp_obj_range_t *self) {
// When computing length, need to take into account step!=1 and step<0.
mp_int_t len = self->stop - self->start + self->step;
if (self->step > 0) {
len -= 1;
} else {
len += 1;
}
len = len / self->step;
if (len < 0) {
len = 0;
}
return len;
}
STATIC mp_obj_t range_unary_op(int op, mp_obj_t self_in) {
mp_obj_range_t *self = self_in;
mp_int_t len = range_len(self);
switch (op) {
case MP_UNARY_OP_BOOL: return MP_BOOL(len > 0);
case MP_UNARY_OP_LEN: return MP_OBJ_NEW_SMALL_INT(len);
default: return MP_OBJ_NULL; // op not supported
}
}
STATIC mp_obj_t range_subscr(mp_obj_t self_in, mp_obj_t index, mp_obj_t value) {
if (value == MP_OBJ_SENTINEL) {
// load
mp_obj_range_t *self = self_in;
mp_int_t len = range_len(self);
#if MICROPY_PY_BUILTINS_SLICE
if (MP_OBJ_IS_TYPE(index, &mp_type_slice)) {
mp_bound_slice_t slice;
mp_seq_get_fast_slice_indexes(len, index, &slice);
mp_obj_range_t *o = m_new_obj(mp_obj_range_t);
o->base.type = &mp_type_range;
o->start = slice.start;
o->stop = slice.stop;
o->step = slice.step;
return o;
}
#endif
uint index_val = mp_get_index(self->base.type, len, index, false);
return MP_OBJ_NEW_SMALL_INT(self->start + index_val * self->step);
} else {
return MP_OBJ_NULL; // op not supported
}
}
STATIC mp_obj_t range_getiter(mp_obj_t o_in) {
mp_obj_range_t *o = o_in;
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) {
const mp_obj_type_t mp_type_range = {
{ &mp_type_type },
.name = MP_QSTR_range,
.print = range_print,
.make_new = range_make_new,
.unary_op = range_unary_op,
.subscr = range_subscr,
.getiter = range_getiter,
};
# test builtin range type
# print
print(range(4))
# len
print(len(range(0)))
print(len(range(4)))
print(len(range(1, 4)))
print(len(range(1, 4, 2)))
print(len(range(1, 4, -1)))
print(len(range(4, 1, -1)))
print(len(range(4, 1, -2)))
# subscr
print(range(4)[0])
print(range(4)[1])
print(range(4)[-1])
# slice
print(range(4)[0:])
print(range(4)[1:])
print(range(4)[1:2])
print(range(4)[1:3])
print(range(4)[1::2])
print(range(4)[1:-2:2])
......@@ -18,8 +18,8 @@ for b in reversed(b'1234'):
print(b)
# range
#for i in reversed(range(3)):
# print(i)
for i in reversed(range(3)):
print(i)
# user object
class A:
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment