Setitem method for ArraySequence

Author: MarcCote

nibabel/streamlines/array_sequence.py

@@ -434,12 +434,26 @@ def __setitem__(self, idx, elements):
             raise TypeError("Index must be either an int, a slice, a list of int"
                             " or a ndarray of bool! Not " + str(type(idx)))
 
-        if len(lengths) != elements.total_nb_rows:
-            msg = "Trying to set {} sequences with {} sequences."
-            raise TypeError(msg.format(len(lengths), elements.total_nb_rows))
+        if is_array_sequence(elements):
+            if len(lengths) != len(elements):
+                msg = "Trying to set {} sequences with {} sequences."
+                raise ValueError(msg.format(len(lengths), len(elements)))
+
+            if sum(lengths) != elements.total_nb_rows:
+                msg = "Trying to set {} points with {} points."
+                raise ValueError(msg.format(sum(lengths), elements.total_nb_rows))
+
+            for o1, l1, o2, l2 in zip(offsets, lengths, elements._offsets, elements._lengths):
+                data[o1:o1 + l1] = elements._data[o2:o2 + l2]
+
+        elif isinstance(elements, numbers.Number):
+            for o1, l1 in zip(offsets, lengths):
+                data[o1:o1 + l1] = elements
+
+        else:  # Try to iterate over it.
+            for o1, l1, data in zip(offsets, lengths, elements):
+                data[o1:o1 + l1] = data
 
-        for o1, l1, o2, l2 in zip(offsets, lengths, elements._offsets, elements._lengths):
-            data[o1:o1 + l1] = elements._data[o2:o2 + l2]
 
     def _op(self, op, value=None, inplace=False):
         """ Applies some operator to this arraysequence.

nibabel/streamlines/tests/test_array_sequence.py

@@ -375,7 +375,6 @@ def _test_binary(op, arrseq, scalars, seqs, inplace=False):
         # Restore flags.
         np.seterr(**flags)
 
-
     def test_arraysequence_setitem(self):
         # Set one item
         seq = SEQ_DATA['seq'] * 0
@@ -384,51 +383,35 @@ def test_arraysequence_setitem(self):
 
         check_arr_seq(seq, SEQ_DATA['seq'])
 
-        # Get all items using indexing (creates a view).
-        indices = list(range(len(SEQ_DATA['seq'])))
-        seq_view = SEQ_DATA['seq'][indices]
-        check_arr_seq_view(seq_view, SEQ_DATA['seq'])
-        # We took all elements so the view should match the original.
-        check_arr_seq(seq_view, SEQ_DATA['seq'])
-
-        # Get multiple items using ndarray of dtype integer.
-        for dtype in [np.int8, np.int16, np.int32, np.int64]:
-            seq_view = SEQ_DATA['seq'][np.array(indices, dtype=dtype)]
-            check_arr_seq_view(seq_view, SEQ_DATA['seq'])
-            # We took all elements so the view should match the original.
-            check_arr_seq(seq_view, SEQ_DATA['seq'])
-
-        # Get multiple items out of order (creates a view).
-        SEQ_DATA['rng'].shuffle(indices)
-        seq_view = SEQ_DATA['seq'][indices]
-        check_arr_seq_view(seq_view, SEQ_DATA['seq'])
-        check_arr_seq(seq_view, [SEQ_DATA['data'][i] for i in indices])
+        # Setitem with a scalar.
+        seq = SEQ_DATA['seq'].copy()
+        seq[:] = 0
+        assert_true(seq._data.sum() == 0)
 
-        # Get slice (this will create a view).
-        seq_view = SEQ_DATA['seq'][::2]
-        check_arr_seq_view(seq_view, SEQ_DATA['seq'])
-        check_arr_seq(seq_view, SEQ_DATA['data'][::2])
+        # Setitem with a list of ndarray.
+        seq = SEQ_DATA['seq'].copy()
+        for i, data in enumerate(SEQ_DATA['data']):
+            seq[i] = data
+        check_arr_seq(seq, SEQ_DATA['data'])
 
-        # Use advanced indexing with ndarray of data type bool.
-        selection = np.array([False, True, True, False, True])
-        seq_view = SEQ_DATA['seq'][selection]
-        check_arr_seq_view(seq_view, SEQ_DATA['seq'])
-        check_arr_seq(seq_view,
-                      [SEQ_DATA['data'][i]
-                       for i, keep in enumerate(selection) if keep])
+        # Setting a slice using another slice.
+        seq = ArraySequence(np.arange(900).reshape((50,6,3)))
+        seq[0:4] = seq[5:9]
+        check_arr_seq(seq[0:4], seq[5:9])
 
-        # Test invalid indexing
-        assert_raises(TypeError, SEQ_DATA['seq'].__getitem__, 'abc')
+        # Setting a slice using another slice with more sequences.
+        seq = ArraySequence(np.arange(900).reshape((50,6,3)))
+        assert_raises(ValueError, seq.__setitem__, slice(0, 4), seq[5:10])
 
-        # Get specific columns.
-        seq_view = SEQ_DATA['seq'][:, 2]
-        check_arr_seq_view(seq_view, SEQ_DATA['seq'])
-        check_arr_seq(seq_view, [d[:, 2] for d in SEQ_DATA['data']])
+        # Setitem between array sequences with different amount of points.
+        seq1 = ArraySequence(np.arange(10).reshape(5, 2))
+        seq2 = ArraySequence(np.arange(15).reshape(5, 3))
+        assert_raises(ValueError, seq1.__setitem__, slice(0, 5), seq2)
 
-        # Combining multiple slicing and indexing operations.
-        seq_view = SEQ_DATA['seq'][::-2][:, 2]
-        check_arr_seq_view(seq_view, SEQ_DATA['seq'])
-        check_arr_seq(seq_view, [d[:, 2] for d in SEQ_DATA['data'][::-2]])
+        # Setitem between array sequences with different common shape.
+        seq1 = ArraySequence(np.arange(12).reshape(2, 2, 3))
+        seq2 = ArraySequence(np.arange(8).reshape(2, 2, 2))
+        assert_raises(ValueError, seq1.__setitem__, slice(0, 2), seq2)
 
     def test_arraysequence_repr(self):
         # Test that calling repr on a ArraySequence object is not falling.