Commit 3eb532e9 authored by Paul Sokolovsky's avatar Paul Sokolovsky
Browse files

extmod/modbtree: Implement __contains__ operation.

parent 8766bc02
...@@ -31,6 +31,7 @@ ...@@ -31,6 +31,7 @@
#include "py/nlr.h" #include "py/nlr.h"
#include "py/runtime.h" #include "py/runtime.h"
#include "py/runtime0.h"
#include "py/stream.h" #include "py/stream.h"
#if MICROPY_PY_BTREE #if MICROPY_PY_BTREE
...@@ -292,6 +293,24 @@ STATIC mp_obj_t btree_subscr(mp_obj_t self_in, mp_obj_t index, mp_obj_t value) { ...@@ -292,6 +293,24 @@ STATIC mp_obj_t btree_subscr(mp_obj_t self_in, mp_obj_t index, mp_obj_t value) {
} }
} }
STATIC mp_obj_t btree_binary_op(mp_uint_t op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
mp_obj_btree_t *self = MP_OBJ_TO_PTR(lhs_in);
switch (op) {
case MP_BINARY_OP_IN: {
mp_uint_t v;
DBT key, val;
key.data = (void*)mp_obj_str_get_data(rhs_in, &v);
key.size = v;
int res = __bt_get(self->db, &key, &val, 0);
CHECK_ERROR(res);
return mp_obj_new_bool(res != RET_SPECIAL);
}
default:
// op not supported
return MP_OBJ_NULL;
}
}
STATIC const mp_rom_map_elem_t btree_locals_dict_table[] = { STATIC const mp_rom_map_elem_t btree_locals_dict_table[] = {
{ MP_ROM_QSTR(MP_QSTR_close), MP_ROM_PTR(&btree_close_obj) }, { MP_ROM_QSTR(MP_QSTR_close), MP_ROM_PTR(&btree_close_obj) },
{ MP_ROM_QSTR(MP_QSTR_get), MP_ROM_PTR(&btree_get_obj) }, { MP_ROM_QSTR(MP_QSTR_get), MP_ROM_PTR(&btree_get_obj) },
...@@ -311,6 +330,7 @@ STATIC const mp_obj_type_t btree_type = { ...@@ -311,6 +330,7 @@ STATIC const mp_obj_type_t btree_type = {
.print = btree_print, .print = btree_print,
.getiter = btree_getiter, .getiter = btree_getiter,
.iternext = btree_iternext, .iternext = btree_iternext,
.binary_op = btree_binary_op,
.subscr = btree_subscr, .subscr = btree_subscr,
.locals_dict = (void*)&btree_locals_dict, .locals_dict = (void*)&btree_locals_dict,
}; };
......
...@@ -62,5 +62,8 @@ print(list(db.values())) ...@@ -62,5 +62,8 @@ print(list(db.values()))
for k in db: for k in db:
print(k) print(k)
print("foo1", "foo1" in db)
print("foo2", "foo2" in db)
db.close() db.close()
f.close() f.close()
...@@ -30,3 +30,5 @@ KeyError ...@@ -30,3 +30,5 @@ KeyError
b'bar1' b'bar1'
b'foo1' b'foo1'
b'foo3' b'foo3'
foo1 True
foo2 False
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