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+
8596STATIC 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+
106167STATIC 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) {
111172const 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};
0 commit comments