added test for __setitem__ using an ndarray value

gdementen · gdementen · commit 09419de82664 · 2015-01-28T12:09:27.000+01:00
diff --git a/larray/tests/test_la.py b/larray/tests/test_la.py
@@ -647,6 +647,28 @@ def test_setitem_larray(self):
         raw[[1, 5, 9]] = raw[[2, 7, 3]] + 27.0
         self._assert_equal_raw(la, raw)
 
+    def test_setitem_ndarray(self):
+        """
+        tests LArray.__setitem__(key, value) where value is a raw ndarray.
+        In that case, value.shape is more restricted as we rely on
+        numpy broadcasting.
+        """
+        # a) value has exactly the same shape as the target slice
+        la = self.larray.copy()
+        raw = self.array.copy()
+        value = raw[[1, 5, 9]] + 25.0
+        la['1,5,9'] = value
+        raw[[1, 5, 9]] = value
+        self._assert_equal_raw(la, raw)
+
+        # b) value has the same axes than target but one has length 1
+        la = self.larray.copy()
+        raw = self.array.copy()
+        value = np.sum(raw[[1, 5, 9]], axis=1, keepdims=True)
+        la['1,5,9'] = value
+        raw[[1, 5, 9]] = value
+        self._assert_equal_raw(la, raw)
+
     def test_set(self):
         age, geo, sex, lipro = self.larray.axes
 
