extmod/modbtree: Use buffer protocol for keys/values.

This changes the btree implementation to use the buffer protocol for
reading key/values in all methods.  `str` and `bytes` objects are not the
only bytes-like objects that could be used.

Documentation and tests are also updated.

Addresses issue #8748.

Signed-off-by: David Lechner <david@pybricks.com>
This commit is contained in:
David Lechner 2022-06-10 12:28:56 -05:00 committed by Damien George
parent c118b5d0e4
commit a565811f23
3 changed files with 26 additions and 15 deletions

View File

@ -11,8 +11,9 @@ value, a database also supports efficient ordered range scans (retrieval
of values with the keys in a given range). On the application interface
side, BTree database work as close a possible to a way standard `dict`
type works, one notable difference is that both keys and values must
be `bytes` objects (so, if you want to store objects of other types, you
need to serialize them to `bytes` first).
be `bytes`-like objects (so, if you want to store objects of other types, you
need to first serialize them to `str` or `bytes` or another type that supports
the buffer protocol).
The module is based on the well-known BerkelyDB library, version 1.xx.

View File

@ -76,6 +76,13 @@ STATIC mp_obj_btree_t *btree_new(DB *db, mp_obj_t stream) {
return o;
}
STATIC void buf_to_dbt(mp_obj_t obj, DBT *dbt) {
mp_buffer_info_t bufinfo;
mp_get_buffer_raise(obj, &bufinfo, MP_BUFFER_READ);
dbt->data = bufinfo.buf;
dbt->size = bufinfo.len;
}
STATIC void btree_print(const mp_print_t *print, mp_obj_t self_in, mp_print_kind_t kind) {
(void)kind;
mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
@ -98,8 +105,8 @@ STATIC mp_obj_t btree_put(size_t n_args, const mp_obj_t *args) {
(void)n_args;
mp_obj_btree_t *self = MP_OBJ_TO_PTR(args[0]);
DBT key, val;
key.data = (void *)mp_obj_str_get_data(args[1], &key.size);
val.data = (void *)mp_obj_str_get_data(args[2], &val.size);
buf_to_dbt(args[1], &key);
buf_to_dbt(args[2], &val);
return MP_OBJ_NEW_SMALL_INT(__bt_put(self->db, &key, &val, 0));
}
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_put_obj, 3, 4, btree_put);
@ -107,7 +114,7 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_put_obj, 3, 4, btree_put);
STATIC mp_obj_t btree_get(size_t n_args, const mp_obj_t *args) {
mp_obj_btree_t *self = MP_OBJ_TO_PTR(args[0]);
DBT key, val;
key.data = (void *)mp_obj_str_get_data(args[1], &key.size);
buf_to_dbt(args[1], &key);
int res = __bt_get(self->db, &key, &val, 0);
if (res == RET_SPECIAL) {
if (n_args > 2) {
@ -126,7 +133,7 @@ STATIC mp_obj_t btree_seq(size_t n_args, const mp_obj_t *args) {
int flags = MP_OBJ_SMALL_INT_VALUE(args[1]);
DBT key, val;
if (n_args > 2) {
key.data = (void *)mp_obj_str_get_data(args[2], &key.size);
buf_to_dbt(args[2], &key);
}
int res = __bt_seq(self->db, &key, &val, flags);
@ -201,7 +208,7 @@ STATIC mp_obj_t btree_iternext(mp_obj_t self_in) {
if (self->start_key != MP_OBJ_NULL) {
int flags = R_FIRST;
if (self->start_key != mp_const_none) {
key.data = (void *)mp_obj_str_get_data(self->start_key, &key.size);
buf_to_dbt(self->start_key, &key);
flags = R_CURSOR;
} else if (desc) {
flags = R_LAST;
@ -219,7 +226,7 @@ STATIC mp_obj_t btree_iternext(mp_obj_t self_in) {
if (self->end_key != mp_const_none) {
DBT end_key;
end_key.data = (void *)mp_obj_str_get_data(self->end_key, &end_key.size);
buf_to_dbt(self->end_key, &end_key);
BTREE *t = self->db->internal;
int cmp = t->bt_cmp(&key, &end_key);
if (desc) {
@ -254,7 +261,7 @@ STATIC mp_obj_t btree_subscr(mp_obj_t self_in, mp_obj_t index, mp_obj_t value) {
if (value == MP_OBJ_NULL) {
// delete
DBT key;
key.data = (void *)mp_obj_str_get_data(index, &key.size);
buf_to_dbt(index, &key);
int res = __bt_delete(self->db, &key, 0);
if (res == RET_SPECIAL) {
mp_raise_type(&mp_type_KeyError);
@ -264,7 +271,7 @@ STATIC mp_obj_t btree_subscr(mp_obj_t self_in, mp_obj_t index, mp_obj_t value) {
} else if (value == MP_OBJ_SENTINEL) {
// load
DBT key, val;
key.data = (void *)mp_obj_str_get_data(index, &key.size);
buf_to_dbt(index, &key);
int res = __bt_get(self->db, &key, &val, 0);
if (res == RET_SPECIAL) {
mp_raise_type(&mp_type_KeyError);
@ -274,8 +281,8 @@ STATIC mp_obj_t btree_subscr(mp_obj_t self_in, mp_obj_t index, mp_obj_t value) {
} else {
// store
DBT key, val;
key.data = (void *)mp_obj_str_get_data(index, &key.size);
val.data = (void *)mp_obj_str_get_data(value, &val.size);
buf_to_dbt(index, &key);
buf_to_dbt(value, &val);
int res = __bt_put(self->db, &key, &val, 0);
CHECK_ERROR(res);
return mp_const_none;
@ -287,7 +294,7 @@ STATIC mp_obj_t btree_binary_op(mp_binary_op_t op, mp_obj_t lhs_in, mp_obj_t rhs
switch (op) {
case MP_BINARY_OP_CONTAINS: {
DBT key, val;
key.data = (void *)mp_obj_str_get_data(rhs_in, &key.size);
buf_to_dbt(rhs_in, &key);
int res = __bt_get(self->db, &key, &val, 0);
CHECK_ERROR(res);
return mp_obj_new_bool(res != RET_SPECIAL);

View File

@ -10,10 +10,13 @@ except ImportError:
f = uio.BytesIO()
db = btree.open(f, pagesize=512)
mv = memoryview(b"bar1foo1")
db[b"foo3"] = b"bar3"
db[b"foo1"] = b"bar1"
db[b"foo2"] = b"bar2"
db[b"bar1"] = b"foo1"
# any type that implements buffer protocol works for key and value
db["foo2"] = "bar2"
db[mv[:4]] = mv[4:]
dbstr = str(db)
print(dbstr[:7], dbstr[-1:])