From 8478bc38bd94be28d29a57f78bc3b2a4673a180c Mon Sep 17 00:00:00 2001
From: waidhoferj <waidhoferj@gmail.com>
Date: Tue, 14 Mar 2023 09:39:43 -0700
Subject: [PATCH] Preliminary Array Insertion

Fixes #126

`YArray.extend` works with preliminary YTypes once again!

## Changes
- Updated `YArray::insert_multiple_at` to correctly check for preliminary types and gracefully handle errors.

## Testing
- Updated `test_inserts` in test_y_array.py to ensure insertions and extensions work with preliminary types.
---
 src/type_conversions.rs |  2 +-
 src/y_array.rs          | 51 +++++++++++++++++++++++------------------
 tests/test_y_array.py   | 12 ++++++++++
 3 files changed, 42 insertions(+), 23 deletions(-)

diff --git a/src/type_conversions.rs b/src/type_conversions.rs
index fb2ca21..cb8c385 100644
--- a/src/type_conversions.rs
+++ b/src/type_conversions.rs
@@ -278,7 +278,7 @@ impl<'a> Prelim for CompatiblePyType<'a> {
                         let mut y_array = v.borrow_mut();
                         if let SharedType::Prelim(items) = y_array.0.to_owned() {
                             let len = array.len();
-                            YArray::insert_multiple_at(&array, txn, len, items);
+                            YArray::insert_multiple_at(&array, txn, len, items).unwrap();
                         }
                         y_array.0 = SharedType::Integrated(array.clone());
                     }
diff --git a/src/y_array.rs b/src/y_array.rs
index db24979..5f5fb3d 100644
--- a/src/y_array.rs
+++ b/src/y_array.rs
@@ -133,7 +133,7 @@ impl YArray {
         let items = Self::py_iter(items)?;
         match &mut self.0 {
             SharedType::Integrated(array) if array.len() >= index => {
-                Self::insert_multiple_at(array, txn, index, items);
+                Self::insert_multiple_at(array, txn, index, items)?;
                 Ok(())
             }
             SharedType::Prelim(vec) if vec.len() >= index as usize => {
@@ -437,35 +437,42 @@ impl YArray {
         }
     }
 
-    pub fn insert_multiple_at(dst: &Array, txn: &mut Transaction, index: u32, src: Vec<PyObject>) {
-        let mut j = index;
-        let mut i = 0;
+    pub fn insert_multiple_at(
+        dst: &Array,
+        txn: &mut Transaction,
+        index: u32,
+        src: Vec<PyObject>,
+    ) -> PyResult<()> {
+        let mut index = index;
+
         Python::with_gil(|py| {
-            while i < src.len() {
+            let mut iter = src
+                .iter()
+                .map(|element| CompatiblePyType::try_from(element.as_ref(py)))
+                .peekable();
+            while iter.peek().is_some() {
                 let mut anys: Vec<Any> = Vec::default();
-                while i < src.len() {
-                    let converted_item: PyResult<Any> =
-                        CompatiblePyType::try_from(src[i].as_ref(py)).and_then(Any::try_from);
-                    if let Ok(any) = converted_item {
-                        anys.push(any);
-                        i += 1;
-                    } else {
-                        println!("{converted_item:?}");
-                        break;
-                    }
+                while let Some(py_type) =
+                    iter.next_if(|element| !matches!(element, Ok(CompatiblePyType::YType(_))))
+                {
+                    let any = Any::try_from(py_type?)?;
+                    anys.push(any)
                 }
 
                 if !anys.is_empty() {
                     let len = anys.len() as u32;
-                    dst.insert_range(txn, j, anys);
-                    j += len;
-                } else {
-                    let wrapper = PyObjectWrapper(src[i].clone());
-                    dst.insert(txn, j, wrapper);
-                    i += 1;
-                    j += 1;
+                    dst.insert_range(txn, index, anys);
+                    index += len;
+                }
+
+                while let Some(y_type) =
+                    iter.next_if(|element| matches!(element, Ok(CompatiblePyType::YType(_))))
+                {
+                    dst.insert(txn, index, y_type?);
+                    index += 1
                 }
             }
+            Ok(())
         })
     }
 
diff --git a/tests/test_y_array.py b/tests/test_y_array.py
index eae5066..7ff17a0 100644
--- a/tests/test_y_array.py
+++ b/tests/test_y_array.py
@@ -26,6 +26,18 @@ def test_inserts():
 
     assert list(x) == expected
 
+    # Ensure that preliminary types can be inserted
+    integrated_array = d1.get_array("prelim_container")
+    inserted_prelim = YArray(["insert"])
+    extended_prelim = YArray(["extend"])
+
+    with d1.begin_transaction() as txn:
+        integrated_array.insert(txn,0,inserted_prelim)
+        integrated_array.extend(txn, [extended_prelim])
+    values = [list(a) for a in integrated_array]
+    assert values == [["insert"], ["extend"]]
+
+
 
 def test_to_string():
     arr = YArray([7, "awesome", True, ["nested"], {"testing": "dicts"}])