diff --git a/bindings/pyroot/cppyy/CPyCppyy/src/Pythonize.cxx b/bindings/pyroot/cppyy/CPyCppyy/src/Pythonize.cxx index 2dcfc52b29dcd..ff284e15ad3bd 100644 --- a/bindings/pyroot/cppyy/CPyCppyy/src/Pythonize.cxx +++ b/bindings/pyroot/cppyy/CPyCppyy/src/Pythonize.cxx @@ -448,9 +448,20 @@ PyObject* VectorIAdd(PyObject* self, PyObject* args, PyObject* /* kwds */) if (PyObject_CheckBuffer(fi) && !(CPyCppyy_PyText_Check(fi) || PyBytes_Check(fi))) { PyObject* vend = PyObject_CallMethodNoArgs(self, PyStrings::gEnd); if (vend) { - PyObject* result = PyObject_CallMethodObjArgs(self, PyStrings::gInsert, vend, fi, nullptr); + // when __iadd__ is overriden, the operation does not end with + // calling the __iadd__ method, but also assigns the result to the + // lhs of the iadd. For example, performing vec += arr, Python + // first calls our override, and then does vec = vec.iadd(arr). + PyObject *it = PyObject_CallMethodObjArgs(self, PyStrings::gInsert, vend, fi, nullptr); Py_DECREF(vend); - return result; + + if (!it) + return nullptr; + + Py_DECREF(it); + // Assign the result of the __iadd__ override to the std::vector + Py_INCREF(self); + return self; } } } diff --git a/bindings/pyroot/pythonizations/test/stl_vector.py b/bindings/pyroot/pythonizations/test/stl_vector.py index 25142fbeae267..99170082fc442 100644 --- a/bindings/pyroot/pythonizations/test/stl_vector.py +++ b/bindings/pyroot/pythonizations/test/stl_vector.py @@ -5,6 +5,20 @@ 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): """ Tests for the pythonizations of std::vector. @@ -45,6 +59,78 @@ 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/18768 + """ + + # 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.