From abc5eb30d73ef5caf0ea811cc1c8727b081bb04a Mon Sep 17 00:00:00 2001
From: Gregory Neverov <42853258+gneverov@users.noreply.github.com>
Date: Mon, 6 Feb 2023 15:07:26 -0800
Subject: [PATCH] add support for array.extend(iterable)

---
 py/objarray.c | 36 +++++++++++++++++++++---------------
 1 file changed, 21 insertions(+), 15 deletions(-)

diff --git a/py/objarray.c b/py/objarray.c
index a66a73a5ac..763ec15cdb 100644
--- a/py/objarray.c
+++ b/py/objarray.c
@@ -458,26 +458,32 @@ STATIC mp_obj_t array_extend(mp_obj_t self_in, mp_obj_t arg_in) {
 
     // allow to extend by anything that has the buffer protocol (extension to CPython)
     mp_buffer_info_t arg_bufinfo;
-    mp_get_buffer_raise(arg_in, &arg_bufinfo, MP_BUFFER_READ);
+    if (mp_get_buffer(arg_in, &arg_bufinfo, MP_BUFFER_READ)) {
+        size_t sz = mp_binary_get_size('@', self->typecode, NULL);
 
-    size_t sz = mp_binary_get_size('@', self->typecode, NULL);
+        // convert byte count to element count
+        size_t len = arg_bufinfo.len / sz;
 
-    // convert byte count to element count
-    size_t len = arg_bufinfo.len / sz;
+        // make sure we have enough room to extend
+        // TODO: alloc policy; at the moment we go conservative
+        if (self->free < len) {
+            self->items = m_renew(byte, self->items, (self->len + self->free) * sz, (self->len + len) * sz);
+            self->free = 0;
+        } else {
+            self->free -= len;
+        }
 
-    // make sure we have enough room to extend
-    // TODO: alloc policy; at the moment we go conservative
-    if (self->free < len) {
-        self->items = m_renew(byte, self->items, (self->len + self->free) * sz, (self->len + len) * sz);
-        self->free = 0;
+        // extend
+        mp_seq_copy((byte *)self->items + self->len * sz, arg_bufinfo.buf, len * sz, byte);
+        self->len += len;
     } else {
-        self->free -= len;
+        // Otherwise argument must be an iterable of items to append
+        mp_obj_t iterable = mp_getiter(arg_in, NULL);
+        mp_obj_t item;
+        while ((item = mp_iternext(iterable)) != MP_OBJ_STOP_ITERATION) {
+            array_append(self_in, item);
+        }
     }
-
-    // extend
-    mp_seq_copy((byte *)self->items + self->len * sz, arg_bufinfo.buf, len * sz, byte);
-    self->len += len;
-
     return mp_const_none;
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_2(array_extend_obj, array_extend);