djsundog
/
circuitpython
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

extmod/modbtree: Cleverly implement "for key in btree:" syntax. 

I.e. make it work like btree.keys(), while still not using a separate
iterator type.
This commit is contained in:
Paul Sokolovsky 2016-06-23 20:08:37 +03:00
parent d5835261af
commit 2f7ebf16de
1 changed files with 24 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
extmod/modbtree.c
View File

@ -49,6 +49,7 @@ typedef struct _mp_obj_btree_t {
#define FLAG_ITER_VALUES 0x80
#define FLAG_ITER_ITEMS 0xc0
byte flags;
byte next_flags;
} mp_obj_btree_t;
STATIC const mp_obj_type_t btree_type;
@ -66,9 +67,9 @@ STATIC mp_obj_btree_t *btree_new(DB *db) {
mp_obj_btree_t *o = m_new_obj(mp_obj_btree_t);
o->base.type = &btree_type;
o->db = db;
o->start_key = MP_OBJ_NULL;
o->end_key = MP_OBJ_NULL;
o->flags = 0;
o->start_key = mp_const_none;
o->end_key = mp_const_none;
o->next_flags = 0;
return o;
}
@ -132,7 +133,7 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_seq_obj, 2, 4, btree_seq);
STATIC mp_obj_t btree_init_iter(size_t n_args, const mp_obj_t *args, byte type) {
mp_obj_btree_t *self = MP_OBJ_TO_PTR(args[0]);
self->flags = type;
self->next_flags = type;
self->start_key = mp_const_none;
self->end_key = mp_const_none;
if (n_args > 1) {
@ -140,7 +141,7 @@ STATIC mp_obj_t btree_init_iter(size_t n_args, const mp_obj_t *args, byte type)
if (n_args > 2) {
self->end_key = args[2];
if (n_args > 3) {
self->flags |= MP_OBJ_SMALL_INT_VALUE(args[3]);
self->next_flags = type | MP_OBJ_SMALL_INT_VALUE(args[3]);
}
}
}
@ -162,6 +163,23 @@ STATIC mp_obj_t btree_items(size_t n_args, const mp_obj_t *args) {
}
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_items_obj, 1, 4, btree_items);
STATIC mp_obj_t btree_getiter(mp_obj_t self_in) {
mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
if (self->next_flags != 0) {
// If we're called immediately after keys(), values(), or items(),
// use their setup for iteration.
self->flags = self->next_flags;
self->next_flags = 0;
} else {
// Otherwise, iterate over all keys.
self->flags = FLAG_ITER_KEYS;
self->start_key = mp_const_none;
self->end_key = mp_const_none;
}
return self_in;
}
STATIC mp_obj_t btree_iternext(mp_obj_t self_in) {
mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
DBT key, val;
@ -267,7 +285,7 @@ STATIC const mp_obj_type_t btree_type = {
// Save on qstr's, reuse same as for module
.name = MP_QSTR_btree,
.print = btree_print,
.getiter = mp_identity,
.getiter = btree_getiter,
.iternext = btree_iternext,
.subscr = btree_subscr,
.locals_dict = (void*)&btree_locals_dict,