[test] stl vector __iadd__ with array.array

Author: aaronj0

bindings/pyroot/pythonizations/test/stl_vector.py

@@ -4,6 +4,24 @@
 
 import ROOT
 
+# Helper function to assert that the elements of an array object and a std::vector proxy are equal
+def assertVec(vec, arr):
+    # cppyy automatically casts random integers to unicode characters,
+    # so do the same in python so the validation doesn't fail
+    if isinstance(arr, array.array) and arr.typecode in ('b', 'B'):
+        arr = [chr(b) for b in arr]
+
+    tc = unittest.TestCase()
+    # first check lengths match
+    tc.assertEqual(
+        len(vec), len(arr),
+        f"Length mismatch: std::vector is {len(vec)}, array is {len(arr)}"
+    )
+
+    tc.assertSequenceEqual(
+        vec, arr,
+        msg="std::vector and array differ, iadd failed"
+    )
 
 class STL_vector(unittest.TestCase):
     """
@@ -45,6 +63,69 @@ def test_stl_vector_boolean(self):
             self.assertFalse(vector.empty())
             self.assertTrue(bool(vector))
 
+    def test_stl_vector_iadd(self):
+        import array
+        import random
+        """
+        Test that the __iadd__ pythonization of std::vector works as expected.
+        This call dispatches to std::insert
+        https://github.com/root-project/root/issues/
+        """
+
+        # we go over all possible numeric PODs
+        # https://docs.python.org/3/library/array.html
+        entry_type = ['char', 'unsigned char','short', 'unsigned short', 'int',
+                      'unsigned int', 'long', 'unsigned long', 'long long', 'float',
+                      'double']
+        array_type = [ 'b', 'B', 'h', 'H', 'i', 'I', 'l', 'L', 'q', 'Q', 'f', 'd']
+
+        typemap = zip(entry_type, array_type)
+        n = 5
+        for dtype in typemap:
+            vec = ROOT.std.vector[dtype[0]]()
+            self.assertTrue(vec.empty())
+            li = [random.randint(1, 100) for _ in range(n)]
+            arr = array.array(dtype[1], li)
+            vec += arr
+            self.assertFalse(vec.empty())
+            assertVec(vec, arr)
+            vec.pop_back()
+            arr = arr[:-1]
+            assertVec(vec, arr)
+    
+    def test_stl_vector_iadd_2D(self):
+        """
+        Test that the __iadd__ pythonization of std::vector works as expected in 2D
+        """
+        initial = [
+            [1, 2, 3],
+            [4, 5, 6],
+            [7, 8, 9],
+        ]
+
+        vec2d = ROOT.std.vector[ROOT.std.vector[int]](initial)
+        self.assertEqual(
+            len(vec2d), len(initial),
+            f"Initial 2D vector row count wrong ({len(vec2d)} vs {len(initial)})"
+        )
+
+        # verify rows before iadd
+        for idx, (subvec, sublist) in enumerate(zip(vec2d, initial)):
+            with self.subTest(row=idx, phase="before"):
+                assertVec(subvec, sublist)
+
+        vec2d += initial
+        expected = initial + initial
+
+        self.assertEqual(
+            len(vec2d), len(expected),
+            f"2D vector row count after iadd wrong ({len(vec2d)} vs {len(expected)})"
+        )
+
+        for idx, (subvec, sublist) in enumerate(zip(vec2d, expected)):
+            with self.subTest(row=idx, phase="after"):
+                assertVec(subvec, sublist)
+
     def test_tree_with_containers(self):
         """
         Test that the boolean conversion of a std::vector works as expected inside a TTree.