Fix indexing behavior in NumpyTupleDataset

chainer_chemistry/dataset/indexer.py

@@ -167,10 +167,7 @@ def _extract_feature(self, data_index, j):
                                                   self.dataset_length))
                 data_index = numpy.argwhere(data_index).ravel()
 
-            if len(data_index) == 1:
-                return self.extract_feature(data_index[0], j)
-            else:
-                res = [self.extract_feature(i, j) for i in data_index]
+            res = [self.extract_feature(i, j) for i in data_index]
         else:
             return self.extract_feature(data_index, j)
         try:

chainer_chemistry/datasets/numpy_tuple_dataset.py

@@ -35,7 +35,8 @@ def __init__(self, *datasets):
 
     def __getitem__(self, index):
         batches = [dataset[index] for dataset in self._datasets]
-        if isinstance(index, slice):
+        if isinstance(index, slice) or isinstance(index, list) or\
+                isinstance(index, numpy.ndarray):
             length = len(batches[0])
             return [tuple([batch[i] for batch in batches])
                     for i in six.moves.range(length)]

tests/datasets_tests/test_numpy_tuple_dataset.py

@@ -16,6 +16,14 @@ def data():
     return a, b, c
 
 
+@pytest.fixture
+def long_data():
+    a = numpy.array([1, 2, 3, 4])
+    b = numpy.array([4, 5, 6, 7])
+    c = numpy.array([[6, 7, 8], [8, 9, 10], [11, 12, 13], [14, 15, 16]])
+    return a, b, c
+
+
 class TestNumpyTupleDataset(object):
 
     def test_len(self, data):
@@ -47,6 +55,38 @@ def test_get_item_slice_index(self, data, index):
             for a, e in six.moves.zip(tuple_a, tuple_e):
                 numpy.testing.assert_array_equal(a, e)
 
+    @pytest.mark.parametrize('index', [numpy.asarray([2, 0])])
+    def test_get_item_ndarray_index(self, long_data, index):
+        dataset = NumpyTupleDataset(*long_data)
+        actual = dataset[index]
+
+        batches = [d[index] for d in long_data]
+        length = len(batches[0])
+        expect = [tuple([batch[i] for batch in batches])
+                  for i in six.moves.range(length)]
+
+        assert len(actual) == len(expect)
+        for tuple_a, tuple_e in six.moves.zip(actual, expect):
+            assert len(tuple_a) == len(tuple_e)
+            for a, e in six.moves.zip(tuple_a, tuple_e):
+                numpy.testing.assert_array_equal(a, e)
+
+    @pytest.mark.parametrize('index', [[2, 0]])
+    def test_get_item_list_index(self, long_data, index):
+        dataset = NumpyTupleDataset(*long_data)
+        actual = dataset[index]
+
+        batches = [d[index] for d in long_data]
+        length = len(batches[0])
+        expect = [tuple([batch[i] for batch in batches])
+                  for i in six.moves.range(length)]
+
+        assert len(actual) == len(expect)
+        for tuple_a, tuple_e in six.moves.zip(actual, expect):
+            assert len(tuple_a) == len(tuple_e)
+            for a, e in six.moves.zip(tuple_a, tuple_e):
+                numpy.testing.assert_array_equal(a, e)
+
     def test_invalid_datasets(self):
         a = numpy.array([1, 2])
         b = numpy.array([1, 2, 3])