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>

3 files changed, 26 insertions, 15 deletions

diff --git a/docs/library/btree.rst b/docs/library/btree.rst
index c093f970f..9d1dcf111 100644
--- a/docs/library/btree.rst
+++ b/docs/library/btree.rst
@@ -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.
 
diff --git a/extmod/modbtree.c b/extmod/modbtree.c
index 60a5beec6..7a1daacb4 100644
--- a/extmod/modbtree.c
+++ b/extmod/modbtree.c
@@ -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);
diff --git a/tests/extmod/btree1.py b/tests/extmod/btree1.py
index 1f4853eaa..876bce4f8 100644
--- a/tests/extmod/btree1.py
+++ b/tests/extmod/btree1.py
@@ -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:])