diff options
| -rw-r--r-- | extmod/modbtree.c | 20 | ||||
| -rw-r--r-- | tests/extmod/btree1.py | 3 | ||||
| -rw-r--r-- | tests/extmod/btree1.py.exp | 2 |
3 files changed, 25 insertions, 0 deletions
diff --git a/extmod/modbtree.c b/extmod/modbtree.c index 4c58d17d30..3fec52fa54 100644 --- a/extmod/modbtree.c +++ b/extmod/modbtree.c @@ -31,6 +31,7 @@ #include "py/nlr.h" #include "py/runtime.h" +#include "py/runtime0.h" #include "py/stream.h" #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) { } } +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[] = { { MP_ROM_QSTR(MP_QSTR_close), MP_ROM_PTR(&btree_close_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 = { .print = btree_print, .getiter = btree_getiter, .iternext = btree_iternext, + .binary_op = btree_binary_op, .subscr = btree_subscr, .locals_dict = (void*)&btree_locals_dict, }; diff --git a/tests/extmod/btree1.py b/tests/extmod/btree1.py index 54c3dc7947..c96cce92d6 100644 --- a/tests/extmod/btree1.py +++ b/tests/extmod/btree1.py @@ -62,5 +62,8 @@ print(list(db.values())) for k in db: print(k) +print("foo1", "foo1" in db) +print("foo2", "foo2" in db) + db.close() f.close() diff --git a/tests/extmod/btree1.py.exp b/tests/extmod/btree1.py.exp index a266d7acfc..2983a09874 100644 --- a/tests/extmod/btree1.py.exp +++ b/tests/extmod/btree1.py.exp @@ -30,3 +30,5 @@ KeyError b'bar1' b'foo1' b'foo3' +foo1 True +foo2 False |