From ccb6251b59a03f4deb1a33009a65263876c69f7f Mon Sep 17 00:00:00 2001 From: David Lechner Date: Fri, 21 Oct 2022 12:04:56 -0500 Subject: [PATCH] pybricks.common.System.storage: use buffer protocol This updates the storage method to allow any object with the buffer protocol instead of just bytes. Also change the missing arg error to a TypeError while we are touching this. Issue: https://github.com/pybricks/pybricks-micropython/pull/121 --- pybricks/common/pb_type_system.c | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/pybricks/common/pb_type_system.c b/pybricks/common/pb_type_system.c index b90a35ab9..1a63f4c94 100644 --- a/pybricks/common/pb_type_system.c +++ b/pybricks/common/pb_type_system.c @@ -115,7 +115,7 @@ STATIC mp_obj_t pb_type_System_storage(size_t n_args, const mp_obj_t *pos_args, mp_int_t offset = mp_obj_get_int(offset_in); // Handle read. - if (write_in == mp_const_none) { + if (read_in != mp_const_none) { byte *data; mp_uint_t size = mp_obj_get_int(read_in); pb_assert(pbsys_program_load_get_user_data(offset, &data, size)); @@ -123,13 +123,16 @@ STATIC mp_obj_t pb_type_System_storage(size_t n_args, const mp_obj_t *pos_args, } // Handle write. - if (read_in == mp_const_none && !mp_obj_is_str(write_in) && mp_obj_is_str_or_bytes(write_in)) { - mp_obj_str_t *obj = ((mp_obj_str_t *)MP_OBJ_TO_PTR(write_in)); - pbsys_program_load_set_user_data(offset, obj->data, obj->len); + if (write_in != mp_const_none) { + mp_buffer_info_t bufinfo; + mp_get_buffer_raise(write_in, &bufinfo, MP_BUFFER_READ); + + pb_assert(pbsys_program_load_set_user_data(offset, bufinfo.buf, bufinfo.len)); + return mp_const_none; } - mp_raise_ValueError(MP_ERROR_TEXT("Must set either read (int) or write (bytes).")); + mp_raise_TypeError(MP_ERROR_TEXT("Must set either read (int) or write (bytes).")); } STATIC MP_DEFINE_CONST_FUN_OBJ_KW(pb_type_System_storage_obj, 0, pb_type_System_storage);